update followup

main.py

@@ -152,13 +152,14 @@ class ChatRequest(BaseModel):
     conversation_id: Optional[str] = Field(None, description="Unique identifier for the conversation")
     model_id: str = Field(..., description="Identifier for the LLM model to use")
     user_id: str = Field(..., description="Unique identifier for the user")
+    enable_followup: bool = Field(default=False, description="Flag to enable follow-up questions")
 
 async def get_api_key(x_api_key: str = Header(...)) -> str:
     if x_api_key != CHAT_AUTH_KEY:
         raise HTTPException(status_code=403, detail="Invalid API key")
     return x_api_key
 
-async def stream_llm_request(api_key: str, llm_request: Dict[str, str]) -> AsyncGenerator[str, None]:
+async def stream_llm_request(api_key: str, llm_request: Dict[str, str]), endpoint_url:str -> AsyncGenerator[str, None]:
     """
     Make a streaming request to the LLM service.
     """
@@ -166,7 +167,7 @@ async def stream_llm_request(api_key: str, llm_request: Dict[str, str]) -> Async
         async with httpx.AsyncClient() as client:
             async with client.stream(
                 "POST",
-                "https://pvanand-audio-chat.hf.space/llm-agent",
+                endpoint_url,
                 headers={
                     "accept": "text/event-stream",
                     "X-API-Key": api_key,
@@ -186,6 +187,7 @@ async def stream_llm_request(api_key: str, llm_request: Dict[str, str]) -> Async
         logger.error(f"Unexpected error occurred while making LLM request: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Unexpected error occurred while making LLM request: {str(e)}")
 
+
 @app.post("/chat/", response_class=StreamingResponse, tags=["Chat"])
 async def chat(request: ChatRequest, background_tasks: BackgroundTasks, api_key: str = Depends(get_api_key)):
     """
@@ -201,23 +203,35 @@ async def chat(request: ChatRequest, background_tasks: BackgroundTasks, api_key:
 
         # Create RAG prompt
         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"
+        
         # Generate conversation_id if not provided
         conversation_id = request.conversation_id or str(uuid.uuid4())
 
-        # Prepare the request for the LLM service
-        llm_request = {
-            "prompt": request.query,
-            "system_message": rag_prompt,
-            "model_id": request.model_id,
-            "conversation_id": conversation_id,
-            "user_id": request.user_id
-        }
+        if request.enable_followup:
+            # Prepare the request for the LLM service
+            pass
+            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/followup-agent"
+            
+        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}")
         async def response_generator():
             full_response = ""
-            async for chunk in stream_llm_request(api_key, llm_request):
+            async 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}")