Update main.py

main.py

@@ -96,7 +96,6 @@ ModelID = Literal[
     "qwen/qwen-72b-chat",
     "google/gemma-2-27b-it"
 ]
-
 class FollowupQueryModel(BaseModel):
     query: str = Field(..., description="User's query for the followup agent")
     model_id: ModelID = Field(
@@ -105,6 +104,7 @@ class FollowupQueryModel(BaseModel):
     )
     conversation_id: str = Field(default_factory=lambda: str(uuid4()), description="Unique identifier for the conversation")
     user_id: str = Field(..., description="Unique identifier for the user")
+    table_id: Optional[str] = Field(None, description="Optional knowledge base table identifier for the query")
 
     class Config:
         schema_extra = {
@@ -113,18 +113,20 @@ class FollowupQueryModel(BaseModel):
                 "model_id": "openai/gpt-4o-mini",
                 "conversation_id": "123e4567-e89b-12d3-a456-426614174000",
                 "user_id": "user123",
-                "tool_call": "auto"
+                "table_id":"digiyatra"
             }
         }
 
 
-
-async def digiyatra_query_table(query: str, db: Annotated[Any, Depends(get_db_connection)], limit: Optional[int] = 5):
+async def digiyatra_query_table(query: FollowupQueryModel, db: Annotated[Any, Depends(get_db_connection)], limit: Optional[int] = 5):
     """Query the digiyatra table."""
+    if query.table_id is None:
+        query.table_id = "digiyatra"
+        
     response = await query_table(
-        table_id="digiyatra",
-        query=query,
-        user_id="digiyatra", 
+        table_id=query.table_id,
+        query=query.query,
+        user_id=query.user_id, 
         limit=limit
     )
     return response.results['data'][0]['text']
@@ -143,7 +145,7 @@ async def followup_agent(query: FollowupQueryModel, background_tasks: Background
                 {"role": "system", "content": FOLLOWUP_DIGIYATRA_PROMPT}
             ]
 
-        digiyatra_response = await digiyatra_query_table(query.query, db=get_db_connection(), limit=5)
+        digiyatra_response = await digiyatra_query_table(query, db=get_db_connection(), limit=5)
         user_query_with_context = f"{query.query} \n\n FAQ Context for ANSWERING: {digiyatra_response}"
         conversations[query.conversation_id].append({"role": "user", "content": user_query_with_context})
         last_activity[query.conversation_id] = time.time()
@@ -182,6 +184,7 @@ async def followup_agent(query: FollowupQueryModel, background_tasks: Background
         logger.error(f"Error in followup_agent: {str(e)}")
         raise HTTPException(status_code=500, detail="An error occurred while processing the followup agent request.")
 
+       
 from fastapi.middleware.cors import CORSMiddleware
 
 app.add_middleware(