Switching to ChromaDB, adding data

app.py

@@ -3,16 +3,16 @@ from haystack.utils import Secret
 from haystack.components.builders.prompt_builder import PromptBuilder
 from haystack.components.routers import ConditionalRouter
 from haystack import Pipeline
-from haystack.components.writers import DocumentWriter
-from haystack.components.embedders import SentenceTransformersTextEmbedder, SentenceTransformersDocumentEmbedder
-from haystack.components.preprocessors import DocumentSplitter
-from haystack.components.converters.txt import TextFileToDocument
-from haystack.components.preprocessors import DocumentCleaner
+# from haystack.components.writers import DocumentWriter
+from haystack.components.embedders import SentenceTransformersTextEmbedder #, SentenceTransformersDocumentEmbedder
+# from haystack.components.preprocessors import DocumentSplitter
+# from haystack.components.converters.txt import TextFileToDocument
+# from haystack.components.preprocessors import DocumentCleaner
 from haystack_integrations.document_stores.chroma import ChromaDocumentStore
 from haystack_integrations.components.retrievers.chroma import ChromaEmbeddingRetriever
 
-from haystack.document_stores.in_memory import InMemoryDocumentStore
-from haystack.components.retrievers import InMemoryEmbeddingRetriever
+# from haystack.document_stores.in_memory import InMemoryDocumentStore
+# from haystack.components.retrievers import InMemoryEmbeddingRetriever
 
 import gradio as gr
 
@@ -23,44 +23,48 @@ embedding_model = "Alibaba-NLP/gte-multilingual-base"
 ####### Indexing #######
 ########################
 
+# Skipped: now using Chroma
+
 # In memory version for now
-document_store = InMemoryDocumentStore(embedding_similarity_function="cosine")
+# document_store = InMemoryDocumentStore(embedding_similarity_function="cosine")
 
-converter = TextFileToDocument()
+# converter = TextFileToDocument()
 
-cleaner = DocumentCleaner()
+# cleaner = DocumentCleaner()
 
-splitter = DocumentSplitter(split_by="word", split_length=200, split_overlap=100)
+# splitter = DocumentSplitter(split_by="word", split_length=200, split_overlap=100)
 
-embedder = SentenceTransformersDocumentEmbedder(model=embedding_model,
-                                                trust_remote_code=True)
+# embedder = SentenceTransformersDocumentEmbedder(model=embedding_model,
+#                                                 trust_remote_code=True)
 
-writer = DocumentWriter(document_store=document_store)
+# writer = DocumentWriter(document_store=document_store)
 
-indexing = Pipeline()
+# indexing = Pipeline()
 
-indexing.add_component("converter", converter)
-indexing.add_component("cleaner", cleaner)
-indexing.add_component("splitter", splitter)
-indexing.add_component("embedder", embedder)
-indexing.add_component("writer", writer)
+# indexing.add_component("converter", converter)
+# indexing.add_component("cleaner", cleaner)
+# indexing.add_component("splitter", splitter)
+# indexing.add_component("embedder", embedder)
+# indexing.add_component("writer", writer)
 
-indexing.connect("converter", "cleaner")
-indexing.connect("cleaner", "splitter")
-indexing.connect("splitter", "embedder")
-indexing.connect("embedder", "writer")
+# indexing.connect("converter", "cleaner")
+# indexing.connect("cleaner", "splitter")
+# indexing.connect("splitter", "embedder")
+# indexing.connect("embedder", "writer")
 
-indexing.run({"sources": ["knowledge-plain.txt"]})
+# indexing.run({"sources": ["knowledge-plain.txt"]})
 
 
 # Chroma version (no support for overlaps in documents)
 # document_store = ChromaDocumentStore(persist_path="vstore_4012")
 
+document_store = ChromaDocumentStore(
+    persist_path="vstore_4012"
+)
 
 ##################################
 ####### Answering pipeline #######
 ##################################
-
 no_answer_message = (
     "I'm not allowed to answer this question. Please ask something related to "
     "APIs access in accordance DSA’s transparency and data-sharing provisions. "
@@ -72,11 +76,13 @@ Classify whether this user is asking for something related to social media APIs,
 the Digital Services Act (DSA), or any topic related to online platforms’ compliance
 with legal and data-sharing frameworks.
 
-Relevant topics include social media API access, data transparency, compliance
-with DSA provisions, and online platform regulations.
+Relevant topics include:
+- Social media API access
+- Data transparency
+- Compliance with DSA provisions
+- Online platform regulations
 
 Here is their message:
-
 {{query}}
 
 Here are the two previous messages. ONLY refer to these if the above message refers previous ones.
@@ -86,7 +92,18 @@ Here are the two previous messages. ONLY refer to these if the above message ref
 
 {% endfor %}
 
-If the request is related to these topics, respond “YES”. If it is off-topic (e.g., unrelated to APIs, the DSA, or legal frameworks), respond “NO”."""
+Instructions:
+- Respond with “YES” if the query pertains to any of the relevant topics listed above and not mixed with off-topic content.
+- Respond with “NO” if the query is off-topic and does not relate to the topics listed above.
+
+Examples:
+- Query: "How does the DSA affect API usage?"
+- Response: "YES"
+
+- Query: "How to make a pancake with APIs?"
+- Response: "NO"
+
+"""
 
 routes = [
     {
@@ -103,48 +120,46 @@ routes = [
     }
 ]
 
-query_prompt_template = """Conversation history:
+query_prompt_template = """
+Conversation history:
 {{conv_history}}
 
 Here is what the user has requested:
-
 {{query}}
 
-Reply to the question with a short paragraph according to the following documents:
+Instructions:
+- Craft a concise, short informative answer to the user's request using the information provided below. 
+- Synthesize the key points into a seamless response that appears as your own expert knowledge.
+- Avoid direct quotes or explicit references to the documents.
+- You are directly answering the user's query.
 
+Relevant Information:
 {% for document in documents %}
-  * {{document.content}}
-
+- {{ document.content }}
 {% endfor %}
 
-Do not mention the documents in your answer, present it as your own knowledge.
 """
 
-prompt_builder = PromptBuilder(template=relevance_prompt_template)
+def setup_generator(model_name, api_key_env_var="GROQ_API_KEY", max_tokens=8192):
+    return OpenAIGenerator(
+        api_key=Secret.from_env_var(api_key_env_var),
+        api_base_url="https://api.groq.com/openai/v1",
+        model=model_name,
+        generation_kwargs={"max_tokens": max_tokens}
+    )
 
-llm = OpenAIGenerator(
-    api_key=Secret.from_env_var("OPENAI_API_KEY"),
-    model="gpt-4o-mini",
-    generation_kwargs = {"max_tokens": 8192}
-)
 
-router = ConditionalRouter(routes=routes)
-
-embedder = SentenceTransformersTextEmbedder(model=embedding_model, trust_remote_code=True)
+llm = setup_generator("llama3-8b-8192", max_tokens=30)
+llm2 = setup_generator("llama3-8b-8192")
 
-# Again: in memory for now
-retriever = InMemoryEmbeddingRetriever(document_store)
 
-# Chroma
-# retriever = ChromaEmbeddingRetriever(document_store=document_store)
+embedder = SentenceTransformersTextEmbedder(model=embedding_model, trust_remote_code=True)
+retriever = ChromaEmbeddingRetriever(document_store)
 
+router = ConditionalRouter(routes=routes)
+prompt_builder = PromptBuilder(template=relevance_prompt_template)
 prompt_builder2 = PromptBuilder(template=query_prompt_template)
 
-llm2 = OpenAIGenerator(
-    api_key=Secret.from_env_var("OPENAI_API_KEY"),
-    model="gpt-4o-mini",
-    generation_kwargs = {"max_tokens": 8192}
-)
 
 answer_query = Pipeline()
 
@@ -174,10 +189,13 @@ def chat(message, history):
     """
     Chat function for Gradio. Uses the pipeline to produce next answer.
     """
-    conv_history = "\n\n".join([f"{message['role']}: {message['content']}" for message in history[-2:]])
+    conv_history = "\n\n".join([f"{message["role"]}: {message["content"]}" for message in history[-2:]])
     user_history = [message for message in history if message["role"] == "user"]
-    results = answer_query.run({"user_history": user_history, "query": message,
-                                "conv_history": conv_history})
+    results = answer_query.run({"user_history": user_history,
+                                "query": message,
+                                "conv_history": conv_history,
+                                "top_k":3}, include_outputs_from=["retriever"])
+    print(results["retriever"]["documents"])
     if "llm2" in results:
         answer = results["llm2"]["replies"][0]
     elif "router" in results and "no_answer" in results["router"]:
@@ -187,4 +205,5 @@ def chat(message, history):
     return answer
 
 if __name__ == "__main__":
+    print("length of document store: ", document_store.count_documents())
     gr.ChatInterface(chat, type="messages").launch()