Update main.py

main.py

@@ -232,6 +232,53 @@ async def chat(request: ChatRequest, background_tasks: BackgroundTasks, api_key:
 
 
 
+@app.post("/digiyatra-chat/", response_class=StreamingResponse, tags=["Chat"])
+async def chat(request: ChatRequest, background_tasks: BackgroundTasks, api_key: str = Depends(get_api_key)):
+    try:
+        document_list = load_embeddings(request.index_id)
+        search_results = embeddings.search(request.query, 6)
+        context = "\n".join([document_list[idx[0]] for idx in search_results])
+
+        rag_prompt = f"Based on the following context, please answer the user's question:\n\nContext:\n{context}\n\nUser's question: {request.query}\n\nAnswer:"
+        system_prompt = "You are a helpful assistant tasked with providing answers using the context provided"
+        
+        conversation_id = request.conversation_id or str(uuid.uuid4())
+
+        if request.enable_followup:
+            llm_request = {
+                "query": rag_prompt,
+                "model_id": 'openai/gpt-4o-mini',
+                "conversation_id": conversation_id,
+                "user_id": request.user_id
+            }
+            endpoint_url = "https://pvanand-general-chat.hf.space/v2/digiyatra-followup"
+        else:
+            llm_request = {
+                "prompt": rag_prompt,
+                "system_message": system_prompt,
+                "model_id": request.model_id,
+                "conversation_id": conversation_id,
+                "user_id": request.user_id
+            }
+            endpoint_url = "https://pvanand-audio-chat.hf.space/llm-agent"
+        
+        logger.info(f"Starting chat response generation for user: {request.user_id} Full request: {llm_request}")
+        
+        def response_generator():
+            full_response = ""
+            for chunk in stream_llm_request(api_key, llm_request, endpoint_url):
+                full_response += chunk
+                yield chunk
+            logger.info(f"Finished chat response generation for user: {request.user_id} Full response: {full_response}")
+        
+        return StreamingResponse(response_generator(), media_type="text/event-stream")
+
+    except Exception as e:
+        logger.error(f"Error in chat endpoint: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error in chat endpoint: {str(e)}")
+
+
+
 @app.on_event("startup")
 async def startup_event():
     check_and_index_csv_files()