update agent

main.py

@@ -1,23 +1,21 @@
+import os
+import time
+import asyncio
+import logging
+import sqlite3
+import tiktoken
+from uuid import uuid4
+from functools import lru_cache
+from typing import Optional, List, Dict, Literal
 from fastapi import FastAPI, HTTPException, Depends, Security, BackgroundTasks
 from fastapi.security import APIKeyHeader
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, Field
-from typing import Literal, List, Dict
-import os
-from functools import lru_cache
 from openai import OpenAI
-from uuid import uuid4
-import tiktoken
-import sqlite3
-import time
-from datetime import datetime, timedelta
-import asyncio
-import requests
-from prompts import CODING_ASSISTANT_PROMPT, NEWS_ASSISTANT_PROMPT, generate_news_prompt, SEARCH_ASSISTANT_PROMPT, generate_search_prompt
-from fastapi_cache import FastAPICache
-from fastapi_cache.backends.inmemory import InMemoryBackend
-from fastapi_cache.decorator import cache
-import logging
+
+# ============================================================================
+# Configuration and Setup
+# ============================================================================
 
 # Configure logging
 logging.basicConfig(
@@ -30,12 +28,15 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# FastAPI app setup
 app = FastAPI()
 
+# API key configuration
 API_KEY_NAME = "X-API-Key"
 API_KEY = os.environ.get("CHAT_AUTH_KEY", "default_secret_key")
 api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)
 
+# Model definitions
 ModelID = Literal[
     "openai/gpt-4o-mini",
     "meta-llama/llama-3-70b-instruct",
@@ -47,45 +48,34 @@ ModelID = Literal[
     "google/gemma-2-27b-it"
 ]
 
-class QueryModel(BaseModel):
-    user_query: str = Field(..., description="User's coding query")
+# Pydantic models
+class LLMAgentQueryModel(BaseModel):
+    prompt: str = Field(..., description="User's query or prompt")
+    system_message: Optional[str] = Field(None, description="Custom system message for the conversation")
     model_id: ModelID = Field(
-        default="meta-llama/llama-3-70b-instruct",
+        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")
+    conversation_id: Optional[str] = Field(None, description="Unique identifier for the conversation")
     user_id: str = Field(..., description="Unique identifier for the user")
 
     class Config:
         schema_extra = {
             "example": {
-                "user_query": "How do I implement a binary search in Python?",
+                "prompt": "How do I implement a binary search in Python?",
+                "system_message": "You are a helpful coding assistant.",
                 "model_id": "meta-llama/llama-3-70b-instruct",
                 "conversation_id": "123e4567-e89b-12d3-a456-426614174000",
                 "user_id": "user123"
             }
         }
 
-class NewsQueryModel(BaseModel):
-    query: str = Field(..., description="News topic to search for")
-    model_id: ModelID = Field(
-        default="openai/gpt-4o-mini",
-        description="ID of the model to use for response generation"
-    )
-    class Config:
-        schema_extra = {
-            "example": {
-                "query": "Latest developments in AI",
-                "model_id": "openai/gpt-4o-mini"
-            }
-        }
-
+# API key and client setup
 @lru_cache()
 def get_api_keys():
     logger.info("Loading API keys")
     return {
         "OPENROUTER_API_KEY": f"sk-or-v1-{os.environ['OPENROUTER_API_KEY']}",
-        "BRAVE_API_KEY": os.environ['BRAVE_API_KEY']
     }
 
 api_keys = get_api_keys()
@@ -98,53 +88,10 @@ last_activity: Dict[str, float] = {}
 # Token encoding
 encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
 
-def limit_tokens(input_string, token_limit=6000):
-    return encoding.decode(encoding.encode(input_string)[:token_limit])
-
-def calculate_tokens(msgs):
-    return sum(len(encoding.encode(str(m))) for m in msgs)
-
-def chat_with_llama_stream(messages, model="openai/gpt-4o-mini", max_llm_history=4, max_output_tokens=2500):
-    logger.info(f"Starting chat with model: {model}")
-    while calculate_tokens(messages) > (8000 - max_output_tokens):
-        if len(messages) > max_llm_history:
-            messages = [messages[0]] + messages[-max_llm_history:]
-        else:
-            max_llm_history -= 1
-            if max_llm_history < 2:
-                error_message = "Token limit exceeded. Please shorten your input or start a new conversation."
-                logger.error(error_message)
-                raise HTTPException(status_code=400, detail=error_message)
-
-    try:
-        response = or_client.chat.completions.create(
-            model=model,
-            messages=messages,
-            max_tokens=max_output_tokens,
-            stream=True
-        )
-        
-        full_response = ""
-        for chunk in response:
-            if chunk.choices[0].delta.content is not None:
-                content = chunk.choices[0].delta.content
-                full_response += content
-                yield content
-        
-        # After streaming, add the full response to the conversation history
-        messages.append({"role": "assistant", "content": full_response})
-        logger.info("Chat completed successfully")
-    except Exception as e:
-        logger.error(f"Error in model response: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Error in model response: {str(e)}")
-
-async def verify_api_key(api_key: str = Security(api_key_header)):
-    if api_key != API_KEY:
-        logger.warning("Invalid API key used")
-        raise HTTPException(status_code=403, detail="Could not validate credentials")
-    return api_key
+# ============================================================================
+# Database Functions
+# ============================================================================
 
-# SQLite setup
 DB_PATH = '/app/data/conversations.db'
 
 def init_db():
@@ -163,8 +110,6 @@ def init_db():
     conn.close()
     logger.info("Database initialized successfully")
 
-init_db()
-
 def update_db(user_id, conversation_id, message, response):
     logger.info(f"Updating database for conversation: {conversation_id}")
     conn = sqlite3.connect(DB_PATH)
@@ -175,6 +120,65 @@ def update_db(user_id, conversation_id, message, response):
     conn.close()
     logger.info("Database updated successfully")
 
+# ============================================================================
+# Utility Functions
+# ============================================================================
+
+def calculate_tokens(msgs):
+    return sum(len(encoding.encode(str(m))) for m in msgs)
+
+def limit_conversation_history(conversation: List[Dict[str, str]], max_tokens: int = 4000) -> List[Dict[str, str]]:
+    """Limit the conversation history to a maximum number of tokens."""
+    limited_conversation = []
+    current_tokens = 0
+
+    for message in reversed(conversation):
+        message_tokens = calculate_tokens([message])
+        if current_tokens + message_tokens > max_tokens:
+            break
+        limited_conversation.insert(0, message)
+        current_tokens += message_tokens
+
+    return limited_conversation
+
+async def verify_api_key(api_key: str = Security(api_key_header)):
+    if api_key != API_KEY:
+        logger.warning("Invalid API key used")
+        raise HTTPException(status_code=403, detail="Could not validate credentials")
+    return api_key
+
+# ============================================================================
+# LLM Interaction Functions
+# ============================================================================
+
+def chat_with_llama_stream(messages, model="meta-llama/llama-3-70b-instruct", max_output_tokens=2500):
+    logger.info(f"Starting chat with model: {model}")
+    try:
+        response = or_client.chat.completions.create(
+            model=model,
+            messages=messages,
+            max_tokens=max_output_tokens,
+            stream=True
+        )
+        
+        full_response = ""
+        for chunk in response:
+            if chunk.choices[0].delta.content is not None:
+                content = chunk.choices[0].delta.content
+                full_response += content
+                yield content
+        
+        # After streaming, add the full response to the conversation history
+        messages.append({"role": "assistant", "content": full_response})
+        logger.info("Chat completed successfully")
+    except Exception as e:
+        logger.error(f"Error in model response: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error in model response: {str(e)}")
+
+# ============================================================================
+# Background Tasks
+# ============================================================================
+
 async def clear_inactive_conversations():
     while True:
         logger.info("Clearing inactive conversations")
@@ -189,320 +193,45 @@ async def clear_inactive_conversations():
         logger.info(f"Cleared {len(inactive_convos)} inactive conversations")
         await asyncio.sleep(60)  # Check every minute
 
+# ============================================================================
+# FastAPI Events and Endpoints
+# ============================================================================
+
 @app.on_event("startup")
 async def startup_event():
     logger.info("Starting up the application")
-    FastAPICache.init(InMemoryBackend(), prefix="fastapi-cache")
+    init_db()
     asyncio.create_task(clear_inactive_conversations())
 
-@app.post("/coding-assistant")
-async def coding_assistant(query: QueryModel, background_tasks: BackgroundTasks, api_key: str = Depends(verify_api_key)):
-    """
-    Coding assistant endpoint that provides programming help based on user queries.
-    Available models:
-    - meta-llama/llama-3-70b-instruct (default)
-    - anthropic/claude-3.5-sonnet
-    - deepseek/deepseek-coder
-    - anthropic/claude-3-haiku
-    - openai/gpt-3.5-turbo-instruct
-    - qwen/qwen-72b-chat
-    - google/gemma-2-27b-it
-    - openai/gpt-4o-mini
-    Requires API Key authentication via X-API-Key header.
-    """
-    logger.info(f"Received coding assistant query: {query.user_query}")
-    if query.conversation_id not in conversations:
-        conversations[query.conversation_id] = [
-            {"role": "system", "content": "You are a helpful assistant proficient in coding tasks. Help the user in understanding and writing code."}
-        ]
-    
-    conversations[query.conversation_id].append({"role": "user", "content": query.user_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
-        background_tasks.add_task(update_db, query.user_id, query.conversation_id, query.user_query, full_response)
-        logger.info(f"Completed coding assistant response for query: {query.user_query}")
-
-    return StreamingResponse(process_response(), media_type="text/event-stream")
-
-# New functions for news assistant
-
-def internet_search(query, type = "web", num_results=20):
-    logger.info(f"Performing internet search for query: {query}, type: {type}")
-    if type == "web":
-        url = "https://api.search.brave.com/res/v1/web/search"
-    else:
-        url = "https://api.search.brave.com/res/v1/news/search"
-
-    headers = {
-        "Accept": "application/json",
-        "Accept-Encoding": "gzip",
-        "X-Subscription-Token": api_keys["BRAVE_API_KEY"]
-    }
-    params = {"q": query}
-
-    response = requests.get(url, headers=headers, params=params)
-
-    if response.status_code != 200:
-        logger.error(f"Failed to fetch search results. Status code: {response.status_code}")
-        return []
-    
-    if type == "web":
-        search_data = response.json()["web"]["results"]
-    else:
-        search_data = response.json()["results"]
-    processed_results = []
-
-    for item in search_data:
-        if not item.get("extra_snippets"):
-            continue
-        
-        result = {
-            "title": item["title"],
-            "snippet": item["extra_snippets"][0],
-            "last_updated": item.get("age", "")
-        }
-        processed_results.append(result)
-
-    logger.info(f"Retrieved {len(processed_results)} search results")
-    return processed_results[:num_results]
-
-@lru_cache(maxsize=100)
-def cached_internet_search(query: str):
-    logger.info(f"Performing cached internet search for query: {query}")
-    return internet_search(query, type = "news")
-
-
-def analyze_news(query):
-    logger.info(f"Analyzing news for query: {query}")
-    news_data = cached_internet_search(query)
-    
-    if not news_data:
-        logger.error("Failed to fetch news data")
-        return "Failed to fetch news data.", []
-
-    # Prepare the prompt for the AI
-    # Use the imported function to generate the prompt (now includes today's date)
-    prompt = generate_news_prompt(query, news_data)
-
-    messages = [
-        {"role": "system", "content": NEWS_ASSISTANT_PROMPT},
-        {"role": "user", "content": prompt}
-    ]
-
-    logger.info("News analysis completed")
-    return messages
-
-@app.post("/news-assistant")
-async def news_assistant(query: NewsQueryModel, api_key: str = Depends(verify_api_key)):
+@app.post("/llm-agent")
+async def llm_agent(query: LLMAgentQueryModel, background_tasks: BackgroundTasks, api_key: str = Depends(verify_api_key)):
     """
-    News assistant endpoint that provides summaries and analysis of recent news based on user queries.
+    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 news assistant query: {query.query}")
-    messages = analyze_news(query.query)
-    
-    if not messages:
-        logger.error("Failed to fetch news data")
-        raise HTTPException(status_code=500, detail="Failed to fetch news data")
-        
-    def process_response():
-        for content in chat_with_llama_stream(messages, model=query.model_id):
-            yield content
-        logger.info(f"Completed news assistant response for query: {query.query}")
+    logger.info(f"Received LLM agent query: {query.prompt}")
 
-    return StreamingResponse(process_response(), media_type="text/event-stream")
-
-class SearchQueryModel(BaseModel):
-    query: str = Field(..., description="Search query")
-    model_id: ModelID = Field(
-        default="openai/gpt-4o-mini",
-        description="ID of the model to use for response generation"
-    )
-    class Config:
-        schema_extra = {
-            "example": {
-                "query": "What are the latest advancements in quantum computing?",
-                "model_id": "meta-llama/llama-3-70b-instruct"
-            }
-        }
-
-def analyze_search_results(query):
-    logger.info(f"Analyzing search results for query: {query}")
-    search_data = internet_search(query, type="web")
-    
-    if not search_data:
-        logger.error("Failed to fetch search data")
-        return "Failed to fetch search data.", []
-
-    # Prepare the prompt for the AI
-    prompt = generate_search_prompt(query, search_data)
-
-    messages = [
-        {"role": "system", "content": SEARCH_ASSISTANT_PROMPT},
-        {"role": "user", "content": prompt}
-    ]
-
-    logger.info("Search results analysis completed")
-    return messages
-
-@app.post("/search-assistant")
-async def search_assistant(query: SearchQueryModel, api_key: str = Depends(verify_api_key)):
-    """
-    Search assistant endpoint that provides summaries and analysis of web search results based on user queries.
-    Requires API Key authentication via X-API-Key header.
-    """
-    logger.info(f"Received search assistant query: {query.query}")
-    messages = analyze_search_results(query.query)
-    
-    if not messages:
-        logger.error("Failed to fetch search data")
-        raise HTTPException(status_code=500, detail="Failed to fetch search data")
-        
-    def process_response():
-        logger.info(f"Generating response using LLM: {messages}")
-        full_response = ""
-        for content in chat_with_llama_stream(messages, model=query.model_id):
-            full_response+=content
-            yield content
-        logger.info(f"Completed search assistant response for query: {query.query}")
-        logger.info(f"LLM Response: {full_response}")
-
-    return StreamingResponse(process_response(), media_type="text/event-stream")
-
-
-from pydantic import BaseModel, Field
-import yaml
-import json
-from yaml.loader import SafeLoader
-
-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,engaging assistant with the following skills, use them, as necessory. Use <response> tag to provide responses well formatted using markdown format.
-
-If the user request needs further clarification, asnwer to your best of ability in a <response>, further analyze the user request and generate clarifying questions with options <clarification>. Else respond with a helpful answer. After providing response, you can also ask followup questions using the same <clarification> tags in yaml format 
-
-<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>
-
-"""
-
-import re
-
-def parse_followup_response(input_text):
-    # Define patterns for response and clarification
-    response_pattern = re.compile(r'<response>(.*?)<\/response>', re.DOTALL)
-    clarification_pattern = re.compile(r'<clarification>(.*?)<\/clarification>', re.DOTALL)
-
-    # Find all matches for response and clarification
-    response_matches = response_pattern.finditer(input_text)
-    clarification_matches = clarification_pattern.finditer(input_text)
-
-    # Initialize variables to keep track of the position
-    last_end = 0
-    combined_response = ""
-    parsed_clarifications = []
-
-    # Combine responses and capture everything in between
-    for response_match in response_matches:
-        # Capture text before the current response tag
-        combined_response += input_text[last_end:response_match.start()].strip() + "\n"
-        # Add the response content
-        combined_response += response_match.group(1).strip() + "\n"
-        # Update the last end position
-        last_end = response_match.end()
-
-    # Check for clarifications and parse them
-    for clarification_match in clarification_matches:
-        # Capture text before the current clarification tag
-        combined_response += input_text[last_end:clarification_match.start()].strip() + "\n"
-        # Process the clarification block
-        clarification_text = clarification_match.group(1).strip()
-        if clarification_text:
-            # Split by "text:" to separate each question block
-            question_blocks = clarification_text.split("- text:")
-            
-            # Loop through each block and extract the question and its options
-            for block in question_blocks[1:]:
-                # Extract the question using regex (up to the "options:" part)
-                question_match = re.search(r'^(.*?)\s*options:', block, re.DOTALL)
-                if question_match:
-                    question = question_match.group(1).strip()
-                    
-                    # Extract the options using regex
-                    options_match = re.search(r'options:\s*(.*?)$', block, re.DOTALL)
-                    if options_match:
-                        options = [option.strip() for option in options_match.group(1).split('-') if option.strip()]
-                        
-                        # Add the parsed question and options to the list
-                        parsed_clarifications.append({'question': question, 'options': options})
-        # Update the last end position
-        last_end = clarification_match.end()
-
-    # Capture any remaining text after the last tag
-    combined_response += input_text[last_end:].strip()
-
-    return combined_response.strip(), parsed_clarifications
-
-
-@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}")
+    # 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. Provide concise and accurate responses."
         conversations[query.conversation_id] = [
-            {"role": "system", "content": FOLLOWUP_AGENT_PROMPT}
+            {"role": "system", "content": system_message}
         ]
-    
-    conversations[query.conversation_id].append({"role": "user", "content": query.query})
+    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 = conversations[query.conversation_id]
+    limited_conversation = limit_conversation_history(conversations[query.conversation_id])
 
     def process_response():
         full_response = ""
@@ -510,25 +239,18 @@ async def followup_agent(query: FollowupQueryModel, background_tasks: Background
             full_response += content
             yield content
 
-        logger.info(f"LLM RAW response for query: {query.query}: {full_response}")
-        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}")
+
+        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")
 
+# ============================================================================
+# Main Execution
+# ============================================================================
+
 if __name__ == "__main__":
     import uvicorn
-    logger.info("Starting the application")
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=8000)