Update main.py

main.py

@@ -17,7 +17,7 @@ from typing import Optional, Annotated
 from langchain_core.runnables import RunnableConfig
 from langgraph.prebuilt import InjectedState
 from document_rag_router import router as document_rag_router
-from document_rag_router import QueryInput, query_collection, SearchResult
+from document_rag_router import QueryInput, query_collection, SearchResult,db
 from fastapi import HTTPException
 import requests
 from sse_starlette.sse import EventSourceResponse
@@ -136,13 +136,40 @@ model = ChatOpenAI(model="gpt-4o-mini", streaming=True)
 
 # Create a prompt template for formatting
 prompt = ChatPromptTemplate.from_messages([
-    ("system", "You are a helpful AI assistant. Current directory contains: {current_files}"),
+    ("system", "You are a helpful AI assistant. The current collection contains the following files: {collection_files}"),
     ("placeholder", "{messages}"),
 ])
 
-def format_for_model(state):
+async def get_collection_files(collection_id: str, user_id: str) -> str:
+    """Get list of files in the specified collection"""
+    try:
+        # Get the full collection name
+        collection_name = f"{user_id}_{collection_id}"
+        
+        # Open the table and convert to pandas
+        table = db.open_table(collection_name)
+        df = table.to_pandas()
+        
+        # Get unique file names
+        unique_files = df['file_name'].unique()
+        
+        # Join the file names into a string
+        return ", ".join(unique_files)
+    except Exception as e:
+        logging.error(f"Error getting collection files: {str(e)}")
+        return f"Error getting files: {str(e)}"
+
+async def format_for_model(state):
+    # Get collection_id and user_id from the state's configurable
+    config = state.get("configurable", {})
+    collection_id = config.get("collection_id")
+    user_id = config.get("user_id")
+    
+    # Get files in the collection
+    collection_files = await get_collection_files(collection_id, user_id) if collection_id and user_id else "No files available"
+    
     return prompt.invoke({
-        "current_files": get_current_files(),
+        "collection_files": collection_files,
         "messages": state["messages"]
     })