from typing import Callable, Dict, List, Optional

from overrides import override
from autogen.agentchat.contrib.retrieve_user_proxy_agent import RetrieveUserProxyAgent
from autogen.retrieve_utils import get_files_from_dir, split_files_to_chunks
import logging

logger = logging.getLogger(__name__)

try:
    import lancedb
except ImportError as e:
    logging.fatal("lancedb is not installed. Try running 'pip install lancedb'")
    raise e        



db_path = "/tmp/lancedb"

def create_lancedb():

    db = lancedb.connect(db_path)
    data = [
                {"vector": [1.1, 1.2], "id": 11, "documents": "This is a test document spark"},
                {"vector": [0.2, 1.8], "id": 22, "documents": "This is another test document"},
                {"vector": [0.1, 0.3], "id": 3, "documents": "This is a third test document spark"},
                {"vector": [0.5, 0.7], "id": 44, "documents": "This is a fourth test document"},
                {"vector": [2.1, 1.3], "id": 55, "documents": "This is a fifth test document spark"},
                {"vector": [5.1, 8.3], "id": 66, "documents": "This is a sixth test document"},
            ]
    try:
        db.create_table("my_table", data)
    except OSError:
        pass

class LancedbRetrieveUserProxyAgent(RetrieveUserProxyAgent):
    def query_vector_db(
        self,
        query_texts,
        n_results=10,
        search_string="",):

        
        if query_texts:
            vector = [0.1, 0.3]
        db = lancedb.connect(db_path)
        table = db.open_table("my_table")
        query = table.search(vector).where(f"documents LIKE '%{search_string}%'").limit(n_results).to_df()
        data ={"ids": query["id"].tolist(), "documents": query["documents"].tolist()}
        return data




    def retrieve_docs(self, problem: str, n_results: int = 20, search_string: str = ""):
        results = self.query_vector_db(
            query_texts=[problem],
            n_results=n_results,
            search_string=search_string,
        )

        self._results = results
        print("doc_ids: ", results["ids"])

from autogen.agentchat.contrib.retrieve_assistant_agent import RetrieveAssistantAgent
from autogen.agentchat.contrib.retrive_lancedb import LancedbRetrieveUserProxyAgent

from autogen.agentchat.contrib.retrive_lancedb import create_lancedb
autogen.ChatCompletion.start_logging()

# # 1. create an RetrieveAssistantAgent instance named "assistant"
assistant = RetrieveAssistantAgent(
    name="assistant",
    system_message="You are a helpful assistant.",
    llm_config={
        "request_timeout": 600,
        "seed": 42,
        "config_list": config_list,
    },
)


ragragproxyagent = LancedbRetrieveUserProxyAgent(
    name="ragproxyagent",
    human_input_mode="NEVER",
    max_consecutive_auto_reply=2,
    retrieve_config={
        "task": "qa",
        "chunk_token_size": 2000,
        "client": "__",
        "embedding_model": "all-mpnet-base-v2",
    },
)

create_lancedb()


# reset the assistant. Always reset the assistant before starting a new conversation.


code_problem = "How can I use FLAML to perform a classification task and use spark to do parallel training. Train 30 seconds and force cancel jobs if time limit is reached."
ragragproxyagent.initiate_chat(assistant, problem=code_problem)

doc_ids:  [3, 44, 11, 22, 55, 66]
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
[<ipython-input-14-d39f5a89dd2c>](https://localhost:8080/#) in <cell line: 38>()
     36 
     37 code_problem = "How can I use FLAML to perform a classification task and use spark to do parallel training. Train 30 seconds and force cancel jobs if time limit is reached."
---> 38 ragragproxyagent.initiate_chat(assistant, problem=code_problem)
     39 
     40 

2 frames
[/content/autogen/autogen/agentchat/contrib/retrieve_user_proxy_agent.py](https://localhost:8080/#) in _get_context(self, results)
    243             if idx <= _doc_idx:
    244                 continue
--> 245             if results["ids"][0][idx] in self._doc_ids:
    246                 continue
    247             _doc_tokens = num_tokens_from_text(doc, custom_token_count_function=self.custom_token_count_function)

TypeError: 'int' object is not subscriptable

from autogen.agentchat.contrib.retrieve_assistant_agent import RetrieveAssistantAgent
from autogen.agentchat.contrib.retrive_lancedb import LancedbRetrieveUserProxyAgent
from autogen.agentchat.contrib.retrive_lancedb import create_lancedb
autogen.ChatCompletion.start_logging()

# 1. create an RetrieveAssistantAgent instance named "assistant"
assistant = RetrieveAssistantAgent(
    name="assistant",
    system_message="You are a helpful assistant.",
    llm_config={
        "request_timeout": 600,
        "seed": 42,
        "config_list": config_list,
    },
)


ragproxyagent = LancedbRetrieveUserProxyAgent(
    name="ragproxyagent",
    human_input_mode="NEVER",
    max_consecutive_auto_reply=10,
    retrieve_config={
        "task": "qa",
        "docs_path": "https://raw.githubusercontent.com/microsoft/autogen/main/README.md",  # change this to your own path, such as https://raw.githubusercontent.com/microsoft/autogen/main/README.md
        "chunk_token_size": 2000,
        "client": "__",
        "embedding_model": "all-mpnet-base-v2",
    },
)

db = lancedb.connect('/tmp/lancedb')
table = db.create_table("pandas_docs", data=[
    {"vector": embeddings.embed_query("Hello World"), "text": "Hello World", "id": "1"}
], mode="overwrite")
docsearch = LanceDB.from_documents(documents, embeddings, connection=table)