chromadb

app.py

@@ -31,6 +31,7 @@ from google.genai import types
 from story_generator import create_story_prompt_from_pdf, generate_video_from_prompt
 from langchain_huggingface import HuggingFaceEndpoint
 from bigquery_uploader import upload_diagnosis_to_bigquery
+from vector_store import embed_and_store_documents, search_documents
 
 print("✅ All libraries imported successfully.")
 
@@ -53,6 +54,7 @@ print("Performing initial setup...")
 VISION_MODEL, PROCESSOR = load_vision_model()
 KB = KnowledgeBase()
 RETRIEVER = KB  # The retriever is now the KB itself
+embed_and_store_documents() # Initialize and load the vector store
 
 # Initialize ADK components for Connected Mode
 adk_components = initialize_adk(VISION_MODEL, PROCESSOR, RETRIEVER)
@@ -169,17 +171,39 @@ def create_field_mode_ui(user_state):
             report_title = diagnosis
             cleaned_diagnosis = clean_diagnosis_text(diagnosis)
 
-            search_query = "healthy maize" if "healthy" in cleaned_diagnosis.lower() else "phosphorus" if "phosphorus" in cleaned_diagnosis.lower() else "Wetin My Eye See So"
-            remedy = search_bigquery_for_remedy(search_query)
+            # --- Hybrid Search ---
+            # 1. Local Vector Store Search
+            local_remedy_list = search_documents(cleaned_diagnosis)
+            local_remedy = local_remedy_list[0] if local_remedy_list else "No remedy found in local knowledge base."
 
-            final_response = f"## Diagnosis Report            **Condition:**          ### {report_title}          ---         ## Suggested Remedy           {remedy}"
+            # 2. BigQuery Search (as fallback or primary)
+            search_query = "healthy maize" if "healthy" in cleaned_diagnosis.lower() else "phosphorus" if "phosphorus" in cleaned_diagnosis.lower() else "Wetin My Eye See So"
+            cloud_remedy = search_bigquery_for_remedy(search_query)
+
+            final_response = f"""
+            ## Diagnosis Report
+            **Condition:**
+            ### {report_title}
+            ---
+            ## Suggested Remedy (from Local Knowledge)
+            {local_remedy}
+            ---
+            ## Suggested Remedy (from Cloud)
+            {cloud_remedy}
+            """
 
             diagnosis_data = {
                 "ai_diagnosis": report_title,
-                "recommended_action": remedy,
+                "recommended_action": local_remedy, # Prioritizing local remedy for logging
+                "confidence_score": None, # Placeholder
                 "farmer_id": user_state.get("uid"),
-                "farmer_feedback": feedback,
+                "gps_latitude": None, # Placeholder
+                "gps_longitude": None, # Placeholder
                 "crop_type": "Maize",
+                "crop_variety": None, # Placeholder
+                "farmer_feedback": feedback,
+                "treatment_applied": None, # Placeholder
+                "outcome_image_id": None, # Placeholder
             }
             upload_diagnosis_to_bigquery(diagnosis_data)
 

requirements.txt

@@ -11,12 +11,13 @@ langchain-huggingface
 langchain-core
 # For our RAG knowledge base (vector store and embeddings)
 sentence-transformers
-faiss-cpu
+chromadb
+pypdf
 pymupdf
 duckduckgo-search
 langgraph
 google-genai
 google-adk
-pypdf
 google-cloud-bigquery
-requests
+requests
+faiss-cpu

vector_store.py

@@ -0,0 +1,61 @@
+import chromadb
+from sentence_transformers import SentenceTransformer
+import os
+
+# --- Constants ---
+MODEL_NAME = "all-MiniLM-L6-v2"
+COLLECTION_NAME = "aura_mind_knowledge"
+KNOWLEDGE_BASE_DIR = "knowledge_base_data"
+
+# --- Initialize ChromaDB and Model ---
+client = chromadb.PersistentClient(path="chroma_db")
+model = SentenceTransformer(MODEL_NAME)
+collection = client.get_or_create_collection(name=COLLECTION_NAME)
+
+def embed_and_store_documents():
+    """
+    Reads documents from the knowledge base directory, generates embeddings,
+    and stores them in ChromaDB.
+    """
+    if collection.count() > 0:
+        print("✅ Knowledge base is already loaded into ChromaDB.")
+        return
+
+    print("Embedding and storing documents in ChromaDB...")
+    documents = []
+    ids = []
+    for filename in os.listdir(KNOWLEDGE_BASE_DIR):
+        if filename.endswith(".txt"):
+            with open(os.path.join(KNOWLEDGE_BASE_DIR, filename), "r") as f:
+                documents.append(f.read())
+                ids.append(filename)
+
+    if documents:
+        embeddings = model.encode(documents).tolist()
+        collection.add(
+            embeddings=embeddings,
+            documents=documents,
+            ids=ids
+        )
+        print(f"✅ Successfully stored {len(documents)} documents in ChromaDB.")
+
+def search_documents(query: str, n_results: int = 1) -> list:
+    """
+    Searches for relevant documents in ChromaDB based on a query.
+
+    Args:
+        query: The search query.
+        n_results: The number of results to return.
+
+    Returns:
+        A list of relevant documents.
+    """
+    if not query:
+        return []
+
+    query_embedding = model.encode([query]).tolist()
+    results = collection.query(
+        query_embeddings=query_embedding,
+        n_results=n_results,
+    )
+    return results['documents'][0] if results['documents'] else []