Add save_summary and get_summaries endpoints to FastAPI app; refactor create_chroma_db to handle single document input

app.py

@@ -4,6 +4,11 @@ from langgraph.agents.rag_agent.graph import graph as rag_graph
 from fastapi import Request
 from fastapi.middleware.cors import CORSMiddleware
 
+from langchain_core.documents import Document
+
+from utils.create_vectordb import create_chroma_db_and_document,query_chroma_db
+
+
 
 
 
@@ -32,6 +37,52 @@ async def summarize(request: Request):
     notes = data.get("notes")
     return  graph.invoke({"urls": urls, "codes": codes, "notes": notes})
 
+
+@app.post("/save_summary")
+async def save_summary(request: Request):
+    data = await request.json()
+    summary = data.get("summary", "")
+    post_id = data.get("post_id", None)
+    title = data.get("title", "")
+    category = data.get("category", "") 
+    tags = data.get("tags", [])
+    references = data.get("references", [])
+    
+    page_content = f"""
+    Title: {title}
+    Category: {category}
+    Tags: {', '.join(tags)}
+    Summary: {summary}
+    """
+    
+    document = Document(
+        page_content=page_content,
+        id = str(post_id)
+        
+    )
+    
+    is_added = create_chroma_db_and_document(document)
+    
+    if not is_added:
+        return {"error": "Failed to save summary to the database." , "status": "error"}
+    
+    return {"message": "Summary saved successfully." , "status": "success"}
+
+@app.post("/summaries")
+async def get_summaries(request: Request):
+    
+  
+    data = await request.json()
+    print(data)
+    query = data.get("query" , "")
+    
+    print(f"Query received: {query}")
+    results = query_chroma_db(query=query)
+    return results
+    
+    
+    
+    
 @app.post("/chat")
 async def chat(request: Request):
     data = await request.json()

utils/create_vectordb.py

@@ -54,7 +54,7 @@ def split_documents(documents, chunk_size=1000, chunk_overlap=200):
     
     return splits
 
-def create_chroma_db(documents, collection_name="corpus_collection", db_dir=DB_DIR):
+def create_chroma_db_and_document(document, collection_name="corpus_collection", db_dir=DB_DIR):
     """Create a Chroma vector database from documents."""
     # Initialize the Gemini embedding function
     gemini_ef = embedding_functions.GoogleGenerativeAiEmbeddingFunction(
@@ -75,17 +75,25 @@ def create_chroma_db(documents, collection_name="corpus_collection", db_dir=DB_D
             embedding_function=gemini_ef
         )
         print(f"Created new collection: {collection_name}")
+        
+    
+    try:
     
-    # Add documents to collection
-    for i, doc in enumerate(documents):
         collection.add(
-            documents=[doc.page_content],
-            metadatas=[doc.metadata],
-            ids=[f"doc_{i}"]
+            documents = [document.page_content],
+            ids = [document.id]
         )
-    
-    print(f"Added {len(documents)} documents to collection {collection_name}")
-    return collection
+        
+        print("Document added to collection successfully.")
+        return True
+    
+    except Exception as e:
+        print(f"Error adding document to collection: {e}")
+        
+        return False
+        
+        
+ 
 
 def query_chroma_db(query: str, collection_name="corpus_collection", n_results=5, db_dir=DB_DIR):
     """Query the Chroma vector database."""
@@ -136,6 +144,10 @@ def main():
         print(f"Source: {metadata.get('source', 'Unknown')}")
     
     print("\nVector database creation and testing complete!")
+    
+
+
+
 
 if __name__ == "__main__":
     main()