add followup-agent

main.py

@@ -377,6 +377,95 @@ async def search_assistant(query: SearchQueryModel, api_key: str = Depends(verif
 
     return StreamingResponse(process_response(), media_type="text/event-stream")
 
+
+from pydantic import BaseModel, Field
+import yaml
+import json
+
+class FollowupQueryModel(BaseModel):
+    query: str = Field(..., description="User's query for the followup agent")
+    model_id: ModelID = Field(
+        default="openai/gpt-4o-mini",
+        description="ID of the model to use for response generation"
+    )
+    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")
+
+    class Config:
+        schema_extra = {
+            "example": {
+                "query": "How can I improve my productivity?",
+                "model_id": "openai/gpt-4o-mini",
+                "conversation_id": "123e4567-e89b-12d3-a456-426614174000",
+                "user_id": "user123"
+            }
+        }
+
+FOLLOWUP_AGENT_PROMPT = """
+You are a helpful assistant with the following skills, use them, as necessary. If the user request needs further clarification, analyze it and generate clarifying questions with options. Else respond with a helpful answer.  <response>response to user request in markdown</response> <clarification> questions:   - text: [First clarifying question]     options:       - [Option 1]       - [Option 2]       - [Option 3]       - [Option 4 (if needed)]   - text: [Second clarifying question]     options:       - [Option 1]       - [Option 2]       - [Option 3]   # Add more questions as needed    # make sure this section is in valid YAML format </clarification>
+"""
+
+def parse_followup_response(response):
+    response_parts = response.split("<response>")
+    if len(response_parts) > 1:
+        response_content = response_parts[1].split("</response>")[0].strip()
+    else:
+        response_content = ""
+
+    clarification_parts = response.split("<clarification>")
+    if len(clarification_parts) > 1:
+        clarification_yaml = clarification_parts[1].split("</clarification>")[0].strip()
+        try:
+            clarification = yaml.safe_load(clarification_yaml)
+        except yaml.YAMLError:
+            clarification = None
+    else:
+        clarification = None
+
+    return response_content, clarification
+
+@app.post("/followup-agent")
+async def followup_agent(query: FollowupQueryModel, background_tasks: BackgroundTasks, api_key: str = Depends(verify_api_key)):
+    """
+    Followup agent endpoint that provides helpful responses or generates clarifying questions based on user queries.
+    Requires API Key authentication via X-API-Key header.
+    """
+    logger.info(f"Received followup agent query: {query.query}")
+
+    if query.conversation_id not in conversations:
+        conversations[query.conversation_id] = [
+            {"role": "system", "content": FOLLOWUP_AGENT_PROMPT}
+        ]
+    
+    conversations[query.conversation_id].append({"role": "user", "content": query.query})
+    last_activity[query.conversation_id] = time.time()
+    
+    # Limit tokens in the conversation history
+    limited_conversation = conversations[query.conversation_id]
+
+    def process_response():
+        full_response = ""
+        for content in chat_with_llama_stream(limited_conversation, model=query.model_id):
+            full_response += content
+            yield content
+        
+        response_content, clarification = parse_followup_response(full_response)
+        
+        result = {
+            "response": response_content,
+            "clarification": clarification
+        }
+        
+        yield "\n\n" + json.dumps(result)
+        
+        # Add the assistant's response to the conversation history
+        conversations[query.conversation_id].append({"role": "assistant", "content": full_response})
+        
+        background_tasks.add_task(update_db, query.user_id, query.conversation_id, query.query, full_response)
+        logger.info(f"Completed followup agent response for query: {query.query}")
+
+    return StreamingResponse(process_response(), media_type="text/event-stream")
+
 if __name__ == "__main__":
     import uvicorn
     logger.info("Starting the application")