Update llm/agents.py

llm/agents.py

@@ -1,90 +1,96 @@
-import asyncio
-import httpx
 import os
 import json
+import httpx
 
-# Load model config at startup
-with open("llm/model_config.json", "r") as f:
+# Load providers and models
+with open("llm/config.json", "r") as f:
     CONFIG = json.load(f)
 
 PROVIDERS = CONFIG["providers"]
-MODEL_PROVIDER_MAPPING = CONFIG["models"]
-
-async def call_model_api(model: str, prompt: str) -> str:
-    provider_key = MODEL_PROVIDER_MAPPING.get(model)
+MODELS = CONFIG["models"]
+
+# === SYSTEM PROMPTS ===
+STRUCTURED_ASSISTANT_PROMPT = """You are a helpful AI assistant.
+
+- Respond to the user’s message in a structured and professional way.
+- Match the length and complexity of your response to the user's input.
+- If the user's input is simple (e.g., "Hi"), reply politely without overexplaining.
+- If the user's input is complex, give a complete and organized answer.
+- Do not repeat the user's prompt.
+- Be direct, helpful, and clear.
+"""
+
+AGGREGATOR_PROMPT = """You are an AI responsible for combining the outputs of multiple AI assistants.
+
+- Read their answers carefully.
+- Identify the best parts from each.
+- Write a single, coherent, and helpful reply.
+- Do not simply merge texts or repeat everything.
+- Match the depth and tone to the user's original input.
+- Keep it natural and conversational.
+"""
+
+# === CORE FUNCTIONS ===
+async def query_llm(model_name, user_input, role_prompt):
+    provider_key = MODELS.get(model_name)
     if not provider_key:
-        raise ValueError(f"No provider configured for model: {model}")
+        return f"Model '{model_name}' is not supported."
 
     provider = PROVIDERS.get(provider_key)
     if not provider:
-        raise ValueError(f"Provider {provider_key} not found in config")
+        return f"Provider '{provider_key}' is not configured."
 
-    url = provider["url"]
+    endpoint = provider["url"]
     api_key_env = provider["key_env"]
     api_key = os.getenv(api_key_env)
 
     if not api_key:
-        raise ValueError(f"Missing API key for provider {provider_key}")
+        return f"API key for provider '{provider_key}' not found."
 
     headers = {
         "Authorization": f"Bearer {api_key}",
-        "Content-Type": "application/json",
+        "Content-Type": "application/json"
     }
 
-    body = {
-        "model": model,
-        "messages": [{"role": "user", "content": prompt}],
-        "temperature": 0.7,
+    payload = {
+        "model": model_name,
+        "messages": [
+            {"role": "system", "content": role_prompt},
+            {"role": "user", "content": user_input}
+        ]
     }
 
-    async with httpx.AsyncClient(timeout=30) as client:
-        response = await client.post(url, headers=headers, json=body)
-        response.raise_for_status()
-        return response.json()["choices"][0]["message"]["content"]
-
-async def query_llm_agent(name: str, prompt: str, settings: dict) -> str:
-    selected_model = settings.get("models", {}).get(name)
+    try:
+        async with httpx.AsyncClient(timeout=60.0) as client:
+            response = await client.post(endpoint, headers=headers, json=payload)
+            response.raise_for_status()
+            data = response.json()
+            return data["choices"][0]["message"]["content"]
+    except Exception as e:
+        return f"Error: {str(e)}"
 
-    if not selected_model:
-        return f"[{name}] No model selected."
 
-    if selected_model not in MODEL_PROVIDER_MAPPING:
-        return f"[{name}] Model '{selected_model}' is not supported."
+async def query_moa_chain(user_input, settings):
+    """Queries LLM-A, LLM-B, LLM-C, and Aggregator in sequence."""
+    llm_a = settings["models"].get("LLM-A")
+    llm_b = settings["models"].get("LLM-B")
+    llm_c = settings["models"].get("LLM-C")
+    aggregator = settings.get("aggregator")
 
-    try:
-        response = await call_model_api(selected_model, prompt)
-        return f"[{name}] {response}"
-    except Exception as e:
-        return f"[{name}] Error: {str(e)}"
-
-async def query_all_llms(prompt: str, settings: dict) -> list:
-    agents = ["LLM-A", "LLM-B", "LLM-C"]
-    tasks = [query_llm_agent(agent, prompt, settings) for agent in agents]
-    return await asyncio.gather(*tasks)
-
-async def query_aggregator(responses: list, settings: dict) -> str:
-    model = settings.get("aggregator")
-    if not model:
-        return "[Aggregator] No aggregator model selected."
-    if model not in MODEL_PROVIDER_MAPPING:
-        return f"[Aggregator] Model '{model}' is not supported."
-
-    system_prompt = (
-        "You are an aggregator AI. Your task is to read the following responses "
-        "from different AI agents and produce a single, high-quality response.\n\n"
-        + "\n\n".join(responses)
+    # Parallel queries to LLM-A, B, C
+    results = await asyncio.gather(
+        query_llm(llm_a, user_input, STRUCTURED_ASSISTANT_PROMPT),
+        query_llm(llm_b, user_input, STRUCTURED_ASSISTANT_PROMPT),
+        query_llm(llm_c, user_input, STRUCTURED_ASSISTANT_PROMPT)
     )
 
-    try:
-        result = await call_model_api(model, system_prompt)
-        return f"[Aggregator] {result}"
-    except Exception as e:
-        return f"[Aggregator] Error: {str(e)}"
-
-def query_all_llms_sync(prompt: str, settings: dict) -> list:
-    return asyncio.run(query_moa_chain(prompt, settings))
+    # Format outputs to feed aggregator
+    combined_content = (
+        f"[LLM-A] {results[0]}\n\n"
+        f"[LLM-B] {results[1]}\n\n"
+        f"[LLM-C] {results[2]}"
+    )
 
-async def query_moa_chain(prompt: str, settings: dict) -> list:
-    responses = await query_all_llms(prompt, settings)
-    aggregator = await query_aggregator(responses, settings)
-    return responses + [aggregator]
+    # Single query to Aggregator (LLM-D)
+    final_response = await query_llm(aggregator, combined_content, AGGREGATOR_PROMPT)
+    return final_response