chroma db

app.py

@@ -1,12 +1,11 @@
 import gradio as gr
 import os
-from smolagents import Tool, CodeAgent, HfApiModel, GradioUI, DuckDuckGoSearchTool, tool
+from smolagents import Tool, CodeAgent, HfApiModel, stream_to_gradio
 import spaces
 from dotenv import load_dotenv
 import datasets
 from langchain.docstore.document import Document
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.retrievers import BM25Retriever
 import chromadb
 from chromadb.utils import embedding_functions
 
@@ -43,8 +42,8 @@ class RetrieverTool(Tool):
     """
     inputs = {
         "query": {
-            "type": "string",
-            "description": "The query to perform. This should be semantically close to your target documents. Use the affirmative form rather than a question.",
+            "type": "List[string]",
+            "description": "The python list of queries to perform. This should be semantically close to your target documents. Use the affirmative form rather than a question.",
         }
     }
     output_type = "string"
@@ -65,14 +64,14 @@ class RetrieverTool(Tool):
             embedding_function=embedding_func,
             metadata={"hnsw:space": "cosine"},
         )
-        collection.add(
+        collection.upsert(
             documents=[doc.page_content for doc in docs],
             ids=[f"id{i}" for i in range(len(docs))],
         )
         self.collection = collection
 
-    def forward(self, query: str) -> str:
-        assert isinstance(query, str), "Your search query must be a string"
+    def forward(self, query: list[str]) -> str:
+        assert isinstance(query, list[str]), "Your search query must be a string"
         docs = self.collection.query(query, n_results=5)
         retrieved_text = "\nRetrieved documents:\n" + "".join(
             [
@@ -127,4 +126,17 @@ if __name__ == "__main__":
         verbosity_level=10,
     )
 
-    GradioUI(agent).launch()
+    def enter_message(new_message, conversation_history):
+        conversation_history.append(gr.ChatMessage(role="user", content=new_message))
+        yield "", conversation_history
+        for msg in stream_to_gradio(agent, new_message):
+            conversation_history.append(msg)
+            yield "", conversation_history
+
+    with gr.Blocks() as b:
+        chatbot = gr.Chatbot(type="messages", height=1000)
+        textbox = gr.Textbox(lines=3, label="")
+        button = gr.Button("reply")
+        button.click(enter_message, [textbox, chatbot], [textbox, chatbot])
+    b.launch()
+    # GradioUI(agent).launch()