reverted commit..

qchapp · qchapp · commit 6a32fb299712 · 2026-02-18T17:37:32.000+01:00
diff --git a/config.yaml b/config.yaml
@@ -1,42 +1,13 @@
 # AI Agent Model Configuration
 
-# Default/fallback model (used for CLI and initial startup)
-agent_model:
-  name: "gpt-5.1"
-  base_url: null                          # null for default OpenAI endpoint
-  api_key_env: "OPENAI_API_KEY"          # Environment variable containing API key
-
-# Default model for UI dropdown (display_name from available_models)
-default_ui_model: "gpt-5.1"
+# Default config
+# agent_model:
+#   name: "gpt-5.1" # "gpt-4o"                     # Model name
+#   base_url: null                          # null for default OpenAI endpoint
+#   api_key_env: "OPENAI_API_KEY"          # Environment variable containing API key
 
-# Available models for UI dropdown
-available_models:
-  - display_name: "gpt-4o-mini"
-    name: "gpt-4o-mini"
-    base_url: null
-    provider: "OpenAI"
-    api_key_env: "OPENAI_API_KEY"
-  
-  - display_name: "gpt-4o"
-    name: "gpt-4o"
-    base_url: null
-    provider: "OpenAI"
-    api_key_env: "OPENAI_API_KEY"
-  
-  - display_name: "gpt-5-mini"
-    name: "gpt-5-mini"
-    base_url: null
-    provider: "OpenAI"
-    api_key_env: "OPENAI_API_KEY"
-
-  - display_name: "gpt-5.1"
-    name: "gpt-5.1"
-    base_url: null
-    provider: "OpenAI"
-    api_key_env: "OPENAI_API_KEY"
-  
-  - display_name: "GPT-OSS-120B [EPFL]"
-    name: "openai/gpt-oss-120b"
-    base_url: "https://inference-rcp.epfl.ch/v1"
-    provider: "EPFL"
-    api_key_env: "EPFL_API_KEY"
+# Using EPFL's inference server
+agent_model:
+  name: "openai/gpt-oss-120b"
+  base_url: "https://inference.rcp.epfl.ch/v1"
+  api_key_env: "EPFL_API_KEY"           # Set EPFL_API_KEY in .env
diff --git a/src/ai_agent/agent/agent.py b/src/ai_agent/agent/agent.py
@@ -6,12 +6,12 @@
 
 from pydantic_ai import Agent, RunContext
 from pydantic_ai.usage import UsageLimits
-from pydantic_ai.models.openai import OpenAIResponsesModel, OpenAIChatModel
+from pydantic_ai.models.openai import OpenAIResponsesModel
 from pydantic_ai.providers.openai import OpenAIProvider
 from pydantic_ai.messages import BinaryContent
 
 from ai_agent.generator.prompts import get_agent_system_prompt
-from ai_agent.generator.schema import ToolSelection, Conversation, ConversationStatus
+from ai_agent.generator.schema import ToolSelection
 from ai_agent.utils.config import get_config
 from .models import AgentToolSelection, ToolRunLog
 from .tools.repo_info_tool import tool_repo_summary, RepoSummaryInput
@@ -44,10 +44,6 @@
         base_url=agent_model_config.base_url,
         api_key=api_key,
     )
-    openai_model = OpenAIChatModel(
-        model_name=agent_model_config.name,
-        provider=provider,
-    )
 else:
     provider = OpenAIProvider(api_key=api_key)
 
@@ -164,51 +160,38 @@ async def search_alternative(
 
 @agent.tool(retries=2, prepare=cap_prepare)
 @limit_tool_calls("repo_info", cap=12)
-async def repo_info(ctx: RunContext[AgentState], url: str, tool_name: str = None) -> dict:
+async def repo_info(ctx: RunContext[AgentState], url: str) -> dict:
     """
     Fetch a short summary of a GitHub repository.
 
     Non-GitHub URLs are ignored; the tool returns a small dict noting
-    that it was skipped. If a tool_name is provided and the URL is not
-    a GitHub URL, the tool will attempt to look up the GitHub URL from
-    the catalog.
-    
-    Args:
-        url: Repository URL or GitHub owner/repo format
-        tool_name: Optional tool name to look up in catalog if URL is not GitHub
+    that it was skipped.
     """
     norm_url = coerce_github_url_or_none(url)
-    
-    # If URL is not a GitHub URL and tool_name is provided, try catalog lookup
-    if not norm_url and tool_name:
-        log.info(f"Non-GitHub URL provided, tool_name={tool_name}, attempting catalog lookup")
-        # The tool_repo_summary will handle the catalog lookup
-        norm_url = url  # Pass through, tool_repo_summary will handle it
-    elif not norm_url:
+    if not norm_url:
         payload = {
             "tool": "repo_info",
             "url": url,
             "skipped": True,
             "reason": "NON_GITHUB_URL",
-            "hint": "Pass a GitHub repo URL or 'owner/repo' to repo_info(url). Optionally provide tool_name for catalog lookup.",
+            "hint": "Pass a GitHub repo URL or 'owner/repo' to repo_info(url).",
             "timestamp": datetime.now().isoformat()
         }
         ctx.deps.tool_calls.append(payload)
         return {k: v for k, v in payload.items() if k != "tool"}
 
     try:
-        out = await tool_repo_summary(RepoSummaryInput(url=norm_url, tool_name=tool_name))
+        out = await tool_repo_summary(RepoSummaryInput(url=norm_url))
     except Exception as e:
         ctx.deps.tool_calls.append(
-            {"tool": "repo_info", "url": norm_url, "tool_name": tool_name, "error": str(e), "timestamp": datetime.now().isoformat()}
+            {"tool": "repo_info", "url": norm_url, "error": str(e), "timestamp": datetime.now().isoformat()}
         )
         raise
 
     ctx.deps.tool_calls.append(
         {
             "tool": "repo_info",
             "url": norm_url,
-            "tool_name": tool_name,
             "truncated": getattr(out, "truncated", False),
             "timestamp": datetime.now().isoformat()
         }
@@ -261,7 +244,6 @@ def run_agent(
     image_bytes: bytes | None = None,
     model: str | None = None,
     base_url: str | None = None,
-    api_key_env: str | None = None,
     top_k: int | None = None,
     num_choices: int | None = None,
     image_metadata: str | None = None,
@@ -333,19 +315,30 @@ def run_agent(
 
     # When model is provided from UI, base_url comes with it (can be None for OpenAI)
     if model:
-        # Use api_key_env from config if provided, otherwise default to OPENAI_API_KEY
-        key_env_name = api_key_env if api_key_env else "OPENAI_API_KEY"
-        runtime_api_key = os.getenv(key_env_name)
-        if not runtime_api_key:
-            raise ValueError(f"{key_env_name} not found in environment. Cannot use this model.")
-        effective_base_url = base_url  # Can be None for OpenAI
-        log.info(f"✓ Using {key_env_name} for model {effective_model}")
-        log.debug(f"{key_env_name} starts with: {runtime_api_key[:10] if runtime_api_key else 'NONE'}... (len={len(runtime_api_key) if runtime_api_key else 0})")
+        if base_url and "inference.rcp.epfl.ch" in base_url:
+            runtime_api_key = os.getenv("EPFL_API_KEY")
+            if not runtime_api_key:
+                raise ValueError("EPFL_API_KEY not found. Cannot use EPFL models without VPN and API key.")
+            effective_base_url = base_url
+            log.info("✓ Using EPFL_API_KEY for EPFL inference server")
+        else:
+            runtime_api_key = os.getenv("OPENAI_API_KEY")
+            if not runtime_api_key:
+                raise ValueError("OPENAI_API_KEY not found. Cannot use OpenAI models.")
+            effective_base_url = base_url  # None for OpenAI
+            log.info("✓ Using OPENAI_API_KEY for OpenAI endpoint")
     else:
-        # No model override - use config defaults
         effective_base_url = agent_model_config.base_url
-        runtime_api_key = api_key  # Already loaded from config at startup
-        log.info(f"✓ Using API key from config for model {effective_model}")
+        if effective_base_url and "inference.rcp.epfl.ch" in effective_base_url:
+            runtime_api_key = os.getenv("EPFL_API_KEY")
+            if not runtime_api_key:
+                raise ValueError("EPFL_API_KEY not found")
+            log.info("✓ Using EPFL_API_KEY from config")
+        else:
+            runtime_api_key = os.getenv("OPENAI_API_KEY")
+            if not runtime_api_key:
+                raise ValueError("OPENAI_API_KEY not found")
+            log.info("✓ Using OPENAI_API_KEY from config")
 
     # Log runtime configuration
     endpoint_display = effective_base_url if effective_base_url else "api.openai.com"
@@ -369,13 +362,7 @@ def run_agent(
             base_url=effective_base_url,
             api_key=runtime_api_key,
         )
-        
-        # Use OpenAIModel (chat/completions) for custom endpoints, OpenAIResponsesModel for default OpenAI
-        if effective_base_url:
-            log.info("Using OpenAIChatModel (chat/completions API) for custom endpoint")
-            runtime_model = OpenAIChatModel(model_name=effective_model, provider=runtime_provider)
-        else:
-            runtime_model = OpenAIResponsesModel(model_name=effective_model, provider=runtime_provider)
+        runtime_model = OpenAIResponsesModel(model_name=effective_model, provider=runtime_provider)
 
         agent_instance = Agent(
             model=runtime_model,
@@ -429,51 +416,27 @@ def run_agent(
         user_prompt = prompt
 
     # ---- 6) Run the agent --------------------------------------------------
-    try:
-        run_result = agent_instance.run_sync(
-            user_prompt,
-            deps=deps,
-            output_type=ToolSelection,
-            usage_limits=UsageLimits(tool_calls_limit=20),
-        )
-        result = run_result.output
-
-        log.info(f"✅ Agent execution complete - choices returned: {len(result.choices)}")
+    run_result = agent_instance.run_sync(
+        user_prompt,
+        deps=deps,
+        output_type=ToolSelection,
+        usage_limits=UsageLimits(tool_calls_limit=20),
+    )
+    result = run_result.output
 
-        # Log usage (helpful, but may not explicitly expose image-specific counters)
-        if run_result.usage:
-            usage = run_result.usage()
-            log.info(
-                f"📊 Usage: total_tokens={usage.total_tokens}, "
-                f"input_tokens={usage.input_tokens}, output_tokens={usage.output_tokens}"
-            )
+    log.info(f"✅ Agent execution complete - choices returned: {len(result.choices)}")
 
-        # Warn if using non-OpenAI endpoint with images
-        if image_bytes and effective_base_url:
-            log.warning("⚠️  Using custom endpoint - confirm the selected model supports vision.")
+    # Log usage (helpful, but may not explicitly expose image-specific counters)
+    if run_result.usage:
+        usage = run_result.usage()
+        log.info(
+            f"📊 Usage: total_tokens={usage.total_tokens}, "
+            f"request_tokens={usage.request_tokens}, response_tokens={usage.response_tokens}"
+        )
 
-    except Exception as e:
-        # Handle global tool quota limit (UsageLimitExceeded) and other errors gracefully
-        error_msg = str(e)
-        log.warning(f"⚠️  Agent execution encountered an error: {error_msg}")
-        
-        # Check if this is a usage limit error (global tool quota)
-        if "UsageLimitExceeded" in str(type(e).__name__) or "tool_calls_limit" in error_msg.lower():
-            log.warning("Global tool call quota reached - continuing with partial results")
-
-            result = ToolSelection(
-                conversation=Conversation(
-                    status=ConversationStatus.COMPLETE,
-                    context="The agent reached the maximum number of tool calls allowed. Please try a more specific query or break down your request into smaller parts.",
-                    question=None,
-                    options=None
-                ),
-                choices=[],
-                explanation="Tool call limit reached during execution. Try refining your query.",
-                reason=None
-            )
-        else:
-            raise
+    if image_bytes and ("inference.rcp.epfl.ch" in endpoint_display):
+        log.warning("⚠️  Using EPFL inference server - confirm the selected model supports vision on that endpoint.")
+        log.warning("   OpenAI billing/dashboard may not reflect image usage when using a non-OpenAI endpoint.")
 
     # ---- 7) Convert raw tool call records into ToolRunLog objects ----------
     for tc in getattr(deps, "tool_calls", []):
diff --git a/src/ai_agent/ui/components.py b/src/ai_agent/ui/components.py
@@ -10,10 +10,33 @@
 
 from .handlers import respond
 from .visualizations import create_tool_usage_chart, create_tool_timeline, create_disabled_tools_display
-from .utils import get_available_models, get_default_ui_model
 
 log = logging.getLogger("chat_components")
 
+# Model configurations with their inference servers
+MODEL_CONFIGS = {
+    # OpenAI models (default endpoint)
+    "gpt-4o-mini": {"name": "gpt-4o-mini", "base_url": None, "provider": "OpenAI"},
+    "gpt-4o": {"name": "gpt-4o", "base_url": None, "provider": "OpenAI"},
+    "gpt-4-turbo": {"name": "gpt-4-turbo", "base_url": None, "provider": "OpenAI"},
+    
+    # EPFL inference server models
+    "openai/gpt-oss-120b [EPFL]": {
+        "name": "openai/gpt-oss-120b",
+        "base_url": "https://inference-rcp.epfl.ch/v1",
+        "provider": "EPFL"
+    },
+    "mistralai/Mistral-Small-3.2-24B-Instruct-2506 [EPFL]": {
+        "name": "mistralai/Mistral-Small-3.2-24B-Instruct-2506",
+        "base_url": "https://inference.rcp.epfl.ch/v1",
+        "provider": "EPFL"
+    },
+}
+
+def get_model_config(model_display_name: str) -> Dict[str, str]:
+    """Get model configuration from display name."""
+    return MODEL_CONFIGS.get(model_display_name, {"name": model_display_name, "base_url": None, "provider": "Unknown"})
+
 
 def create_chat_interface(doc_index: Dict[str, SoftwareDoc]):
     """
@@ -102,7 +125,7 @@ def create_chat_interface(doc_index: Dict[str, SoftwareDoc]):
         with gr.Row(elem_classes="main-header"):
             gr.HTML("""
                 <div class="logo-container">
-                    <img src="https://imaging-plaza.epfl.ch/logos/imaging_plaza.svg" 
+                    <img src="https://imaging-plaza.epfl.ch/logos/imaging_plaza_white.svg" 
                          alt="Imaging Plaza Logo" 
                          style="height: 48px; width: auto;" />
                     <div>
@@ -115,14 +138,9 @@ def create_chat_interface(doc_index: Dict[str, SoftwareDoc]):
         # Settings section (collapsed by default)
         with gr.Accordion("⚙️ Settings", open=False):
             with gr.Row():
-                # Load models and default from config
-                available_models = get_available_models()
-                model_choices = [m["display_name"] for m in available_models]
-                default_model = get_default_ui_model()
-                
                 model_dropdown = gr.Dropdown(
-                    choices=model_choices,
-                    value=default_model,
+                    choices=list(MODEL_CONFIGS.keys()),
+                    value="gpt-4o-mini",
                     label="Model",
                     info="Select AI model and inference server",
                 )
diff --git a/src/ai_agent/ui/handlers.py b/src/ai_agent/ui/handlers.py
@@ -15,7 +15,6 @@
 
 from .state import ChatState, ChatMessage
 from .formatters import format_tool_card
-from .utils import get_model_config
 
 log = logging.getLogger("chat_handlers")
 
@@ -202,13 +201,13 @@ def respond(
     # Parse model configuration if provided
     model_name = None
     base_url_override = None  # Use different variable name
-    api_key_env = None
     if model:
+        # Import here to avoid circular dependency
+        from ai_agent.ui.components import get_model_config
         model_config = get_model_config(model)
         model_name = model_config.get("name")
         base_url_override = model_config.get("base_url")  # Can be None for OpenAI
-        api_key_env = model_config.get("api_key_env", "OPENAI_API_KEY")
-        log.info(f"Model config: {model} -> name={model_name}, base_url={base_url_override}, api_key_env={api_key_env}")
+        log.info(f"Model config: {model} -> name={model_name}, base_url={base_url_override}")
 
     effective_paths = file_paths or (state.last_files or [])
 
@@ -229,7 +228,7 @@ def respond(
             conversation_history=state.conversation_history,
             model=model_name,
             base_url=base_url_override if model else None,  # Only override if model selected
-            api_key_env=api_key_env, 
+            top_k=top_k,
             num_choices=num_choices,
         )
     except ValueError as e:
diff --git a/src/ai_agent/ui/utils.py b/src/ai_agent/ui/utils.py
