Update main.py

main.py

@@ -277,6 +277,51 @@ async def llm_agent(query: LLMAgentQueryModel, background_tasks: BackgroundTasks
 
     return StreamingResponse(process_response(), media_type="text/event-stream")
 
+@app.post("/v2/llm-agent")
+async def llm_agent_v2(query: LLMAgentQueryModel, background_tasks: BackgroundTasks, api_key: str = Depends(verify_api_key)):
+    """
+    LLM agent endpoint that provides responses based on user queries, maintaining conversation history.
+    Accepts custom system messages and allows selection of different models.
+    Requires API Key authentication via X-API-Key header.
+    """
+    logger.info(f"Received LLM agent query: {query.prompt}")
+
+    # Generate a new conversation ID if not provided
+    if not query.conversation_id:
+        query.conversation_id = str(uuid4())
+
+    # Initialize or retrieve conversation history
+    if query.conversation_id not in conversations:
+        system_message = query.system_message or "You are a helpful assistant."
+        conversations[query.conversation_id] = [
+            {"role": "system", "content": system_message}
+        ]
+    elif query.system_message:
+        # Update system message if provided
+        conversations[query.conversation_id][0] = {"role": "system", "content": query.system_message}
+
+    # Add user's prompt to conversation history
+    conversations[query.conversation_id].append({"role": "user", "content": query.prompt})
+    last_activity[query.conversation_id] = time.time()
+
+    # Limit tokens in the conversation history
+    limited_conversation = limit_conversation_history(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 json.dumps({"type": "response","content": content}) + "\n"
+
+        # 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.prompt, full_response)
+        logger.info(f"Completed LLM agent response for query: {query.prompt}")
+
+    return StreamingResponse(process_response(), media_type="text/event-stream")
+
+
 import edge_tts
 import io