fixed config usage and logo

qchapp · qchapp · commit 4db1f38207a3 · 2026-02-18T17:33:11.000+01:00
diff --git a/config.yaml b/config.yaml
@@ -1,13 +1,42 @@
 # AI Agent Model Configuration
 
-# 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
-
-# Using EPFL's inference server
+# Default/fallback model (used for CLI and initial startup)
 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
+  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"
+
+# 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"
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
+from pydantic_ai.models.openai import OpenAIResponsesModel, OpenAIChatModel
 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
+from ai_agent.generator.schema import ToolSelection, Conversation, ConversationStatus
 from ai_agent.utils.config import get_config
 from .models import AgentToolSelection, ToolRunLog
 from .tools.repo_info_tool import tool_repo_summary, RepoSummaryInput
@@ -44,6 +44,10 @@
         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)
 
@@ -160,38 +164,51 @@ 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) -> dict:
+async def repo_info(ctx: RunContext[AgentState], url: str, tool_name: str = None) -> 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.
+    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
     """
     norm_url = coerce_github_url_or_none(url)
-    if not norm_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:
         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).",
+            "hint": "Pass a GitHub repo URL or 'owner/repo' to repo_info(url). Optionally provide tool_name for catalog lookup.",
             "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))
+        out = await tool_repo_summary(RepoSummaryInput(url=norm_url, tool_name=tool_name))
     except Exception as e:
         ctx.deps.tool_calls.append(
-            {"tool": "repo_info", "url": norm_url, "error": str(e), "timestamp": datetime.now().isoformat()}
+            {"tool": "repo_info", "url": norm_url, "tool_name": tool_name, "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()
         }
@@ -244,6 +261,7 @@ 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,
@@ -315,30 +333,19 @@ def run_agent(
 
     # When model is provided from UI, base_url comes with it (can be None for OpenAI)
     if model:
-        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")
+        # 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})")
     else:
+        # No model override - use config defaults
         effective_base_url = agent_model_config.base_url
-        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")
+        runtime_api_key = api_key  # Already loaded from config at startup
+        log.info(f"✓ Using API key from config for model {effective_model}")
 
     # Log runtime configuration
     endpoint_display = effective_base_url if effective_base_url else "api.openai.com"
@@ -362,7 +369,13 @@ def run_agent(
             base_url=effective_base_url,
             api_key=runtime_api_key,
         )
-        runtime_model = OpenAIResponsesModel(model_name=effective_model, provider=runtime_provider)
+        
+        # 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)
 
         agent_instance = Agent(
             model=runtime_model,
@@ -416,27 +429,51 @@ def run_agent(
         user_prompt = prompt
 
     # ---- 6) Run the agent --------------------------------------------------
-    run_result = agent_instance.run_sync(
-        user_prompt,
-        deps=deps,
-        output_type=ToolSelection,
-        usage_limits=UsageLimits(tool_calls_limit=20),
-    )
-    result = run_result.output
+    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)}")
+        log.info(f"✅ Agent execution complete - choices returned: {len(result.choices)}")
 
-    # 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}"
-        )
+        # 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}"
+            )
+
+        # 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.")
 
-    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.")
+    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
 
     # ---- 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,33 +10,10 @@
 
 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]):
     """
@@ -125,7 +102,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_white.svg" 
+                    <img src="https://imaging-plaza.epfl.ch/logos/imaging_plaza.svg" 
                          alt="Imaging Plaza Logo" 
                          style="height: 48px; width: auto;" />
                     <div>
@@ -138,9 +115,14 @@ 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=list(MODEL_CONFIGS.keys()),
-                    value="gpt-4o-mini",
+                    choices=model_choices,
+                    value=default_model,
                     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,6 +15,7 @@
 
 from .state import ChatState, ChatMessage
 from .formatters import format_tool_card
+from .utils import get_model_config
 
 log = logging.getLogger("chat_handlers")
 
@@ -201,13 +202,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
-        log.info(f"Model config: {model} -> name={model_name}, base_url={base_url_override}")
+        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}")
 
     effective_paths = file_paths or (state.last_files or [])
 
@@ -228,7 +229,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
-            top_k=top_k,
+            api_key_env=api_key_env, 
             num_choices=num_choices,
         )
     except ValueError as e:
diff --git a/src/ai_agent/ui/utils.py b/src/ai_agent/ui/utils.py
