Merge pull request #14 from Imaging-Plaza/feature/interface-refactor

Feature/interface refactor

Author: qchapp

CHANGELOG.md

@@ -5,10 +5,34 @@ All notable changes to this project will be documented in this file.
 ## [Unreleased]
 
 ### Added
-- **DeepWiki MCP integration**: Repository info tool now uses DeepWiki MCP server (https://mcp.deepwiki.com/sse) as primary source for GitHub repository documentation. DeepWiki provides fast, pre-indexed documentation access without API rate limits.
-- Automatic fallback to `repocards` library (replacing previous direct GitHub API implementation) when DeepWiki is unavailable or times out, ensuring robust repository information retrieval for both indexed and newly-created repositories.
+- **New chat-based interface** (`ai_agent chat`) with conversational AI assistant
+  - Chatbot component with rich media rendering (images, files, JSON, code blocks)
+  - Inline file upload support for PNG, JPG, WEBP, TIFF, DICOM, NIfTI, CSV, JSON, XML, MP3, MP4
+  - File previews with format-specific icons rendered in chat messages
+  - Tool recommendation cards with detailed metadata (modality, dimensions, license, tags)
+  - Demo execution as conversational flow - assistant asks "Would you like me to run the demo?"
+  - Tool execution traces displayed as collapsible `<details>` sections after responses
+  - Debug sidebar showing conversation state, excluded tools, and preview images
+  - Full conversation context maintained across multi-turn interactions
+  - Affirmative response detection for demo confirmations (yes, sure, ok, etc.)
+- `respond(message, files, state) -> (reply, media, state)` core interface function
+  - Encapsulates all agent logic in testable, UI-independent function
+  - State management via `ChatState` dataclass with serialization
+  - `ChatMessage` dataclass for rich reply composition with markdown, images, files, traces
+- `handlers.py` module with agent response logic
+- `components.py` module for reusable chat UI components
+- `formatters.py` helpers for rich message and media formatting
+- `state.py` chat state models and serialization utilities
+- `visualizations.py` helpers for rendering previews, traces, and visual state
+- `app.py` Gradio app implementing the chat UI
+- **Imaging Plaza branding**: Custom CSS theme with Plaza green colors (#00A991)
+- **Logo integration**: Official Imaging Plaza white logo displayed in header
+- **Redesigned layout**: Reorganized UI with header banner, left chat panel, and right sidebar for files and state
 
 ### Changed
+- CLI now supports `ai_agent chat`
+- **DeepWiki MCP integration**: Repository info tool now uses DeepWiki MCP server (https://mcp.deepwiki.com/sse) as primary source for GitHub repository documentation. DeepWiki provides fast, pre-indexed documentation access without API rate limits.
+- Automatic fallback to `repocards` library (replacing previous direct GitHub API implementation) when DeepWiki is unavailable or times out, ensuring robust repository information retrieval for both indexed and newly-created repositories.
 - Updated `pydantic-ai` dependency to include MCP support via `pydantic-ai[mcp]` extra.
 - Enhanced `RepoSummaryOutput` schema to include `source` field indicating whether data came from "deepwiki" or "repocards".
 - Repository info tool logs now track data source (DeepWiki vs repocards) for observability.
@@ -23,6 +47,9 @@ All notable changes to this project will be documented in this file.
 - **UI State Management Simplified**: Removed complex refine intent detection system. Agent now naturally handles requests for alternatives via conversation history without hard-coded heuristics.
 - **UI Handler Simplified**: Reduced `handle_message()` parameters from 8 to 6, removing `last_task_state`, `last_suggestions_state`, and `excluded_names` state tracking.
 - **Agent-Only Path**: Removed `USE_AGENT` conditional (always uses Pydantic AI agent). Deleted dead code path for non-agent pipeline invocation.
+- **UI redesign**: File upload moved to dedicated right panel for cleaner workflow
+- **Visual hierarchy**: Header with gradient green banner and logo
+- **Button styling**: Primary actions use Imaging Plaza green theme colors
 
 ### Removed
 - **VLMToolSelector**: Deleted unused `generator/generator.py` containing VLMToolSelector class. The pydantic-ai agent handles all tool selection directly.
@@ -32,6 +59,7 @@ All notable changes to this project will be documented in this file.
 - **Legacy Method**: Removed `recommend_and_link()` method from `api/pipeline.py` (~180 lines) - only used by outdated tests, replaced by agent-based approach.
 - **State Variables**: Removed 3 Gradio State objects: `last_task_state`, `last_suggestions_state`, `excluded_names`.
 - **Outdated Tests**: Removed `tests/full_test.py` which only tested the removed `recommend_and_link()` method.
+- CLI no more supports `ai_agent ui` command
 
 ### Fixed
 - **Conversation Context**: Agent now properly maintains conversation history, enabling natural understanding of follow-up requests like "show me alternatives".

pyproject.toml

@@ -6,25 +6,26 @@ readme = "README.md"
 requires-python = ">=3.10"
 
 dependencies = [
-  "faiss-cpu",
-  "numpy",
-  "pydantic>=2",
-  "sentence-transformers",
-  "openai>=1.30.0",
-  "pydantic-ai[mcp]",
-  "requests",
-  "python-dotenv",
+  "faiss-cpu==1.11.0.post1",
+  "numpy==2.2.6",
+  "pydantic==2.11.7",
+  "sentence-transformers==5.1.0",
+  "openai==2.1.0",
+  "pydantic-ai[mcp]==1.0.14",
+  "requests==2.32.4",
+  "python-dotenv==1.1.1",
   "gradio==5.42.0",
-  "json5",
-  "pillow",
-  "nibabel",
-  "tifffile",
-  "pydicom",
-  "imageio",
-  "rdflib",
-  "sparqlwrapper",
-  "repocards",
-  "pyyaml",
+  "json5==0.12.0",
+  "pillow==11.3.0",
+  "nibabel==5.3.2",
+  "tifffile==2025.5.10",
+  "pydicom==3.0.1",
+  "imageio==2.37.0",
+  "rdflib==7.4.0",
+  "sparqlwrapper==2.0.0",
+  "plotly==6.5.0",
+  "repocards==0.1.2",
+  "pyyaml==6.0.2",
 ]
 
 [project.scripts]

src/ai_agent/agent/agent.py

@@ -1,17 +1,18 @@
 from __future__ import annotations
 
 import os, logging
+from datetime import datetime
 from typing import List
 from pydantic_ai import Agent, RunContext
 from pydantic_ai.usage import UsageLimits
 from pydantic_ai.models.openai import OpenAIChatModel
 from pydantic_ai.providers.openai import OpenAIProvider
 
-from generator.prompts import AGENT_SYSTEM_PROMPT
-from generator.schema import ToolSelection
-from api.pipeline import RAGImagingPipeline
-from utils.utils import _best_runnable_link
-from utils.config import get_config
+from ai_agent.generator.prompts import get_agent_system_prompt
+from ai_agent.generator.schema import ToolSelection
+from ai_agent.api.pipeline import RAGImagingPipeline
+from ai_agent.utils.utils import _best_runnable_link
+from ai_agent.utils.config import get_config
 from .models import AgentToolSelection, ToolRunLog
 from .tools.repo_info_tool import tool_repo_summary, RepoSummaryInput
 from .tools.rerank_tool import tool_rerank, RerankInput
@@ -53,7 +54,7 @@
 
 agent = Agent(
     model=openai_model,
-    system_prompt=AGENT_SYSTEM_PROMPT,
+    system_prompt=get_agent_system_prompt(os.getenv("NUM_CHOICES", "3")),
     deps_type=AgentState,
 )
 
@@ -64,16 +65,18 @@
 async def search_tools(ctx: RunContext[AgentState], query: str, excluded: List[str] | None = None, top_k: int = 12, original_formats: List[str] | None = None):
     # Merge explicit excluded param with state's excluded_tools
     all_excluded = list(set((excluded or []) + ctx.deps.excluded_tools))
-    out = tool_search_tools(SearchToolsInput(query=query, excluded=all_excluded, top_k=top_k, original_formats=original_formats or []))
+    # Use override from context if available
+    effective_top_k = ctx.deps.override_top_k if ctx.deps.override_top_k is not None else top_k
+    out = tool_search_tools(SearchToolsInput(query=query, excluded=all_excluded, top_k=effective_top_k, original_formats=original_formats or []))
     payload = [c.model_dump(mode="python") for c in out.candidates]
-    ctx.deps.tool_calls.append({"tool": "search_tools", "query": query, "count": len(payload), "original_formats": original_formats or [], "excluded": all_excluded})
+    ctx.deps.tool_calls.append({"tool": "search_tools", "query": query, "count": len(payload), "original_formats": original_formats or [], "excluded": all_excluded, "timestamp": datetime.now().isoformat()})
     return payload
 
 @agent.tool(retries=2, prepare=cap_prepare)
 @limit_tool_calls("rerank", cap=1)
 async def rerank(ctx: RunContext[AgentState], query: str, candidate_names: List[str], top_k: int = 5):
     out = tool_rerank(RerankInput(query=query, candidate_names=candidate_names, top_k=top_k))
-    ctx.deps.tool_calls.append({"tool": "rerank", "query": query, "used_model": out.used_model, "count": len(out.reranked)})
+    ctx.deps.tool_calls.append({"tool": "rerank", "query": query, "used_model": out.used_model, "count": len(out.reranked), "timestamp": datetime.now().isoformat()})
     return out.model_dump(mode="python")
 
 # @agent.tool(retries=2, prepare=cap_prepare)
@@ -113,7 +116,7 @@ async def repo_info(ctx: RunContext[AgentState], url: str):
             "hint": "Pass a GitHub repo URL or 'owner/repo' to repo_info(url).",
             "original": url,
         }
-        ctx.deps.tool_calls.append({"tool": "repo_info", "url": url, "skipped": True, "reason": "NON_GITHUB_URL"})
+        ctx.deps.tool_calls.append({"tool": "repo_info", "url": url, "skipped": True, "reason": "NON_GITHUB_URL", "timestamp": datetime.now().isoformat()})
         return payload
 
     try:
@@ -122,11 +125,12 @@ async def repo_info(ctx: RunContext[AgentState], url: str):
             "tool": "repo_info",
             "url": norm_url,
             "truncated": out.truncated,
-            "source": out.source
+            "source": out.source,
+            "timestamp": datetime.now().isoformat()
         })
         return out.model_dump(mode="python")
     except Exception as e:
-        ctx.deps.tool_calls.append({"tool": "repo_info", "url": norm_url, "error": str(e)})
+        ctx.deps.tool_calls.append({"tool": "repo_info", "url": norm_url, "error": str(e), "timestamp": datetime.now().isoformat()})
         return {
             "invalid": True,
             "reason": "FETCH_FAILED",
@@ -146,14 +150,23 @@ async def resolve_demo_link(ctx: RunContext[AgentState], tool_name: str):
             link = _best_runnable_link(doc)
     except Exception:
         link = None
-    ctx.deps.tool_calls.append({"tool": "resolve_demo_link", "tool_name": tool_name, "demo_link": link})
+    ctx.deps.tool_calls.append({"tool": "resolve_demo_link", "tool_name": tool_name, "demo_link": link, "timestamp": datetime.now().isoformat()})
     return {"tool_name": tool_name, "demo_link": link}
 
 # Runner wrapper ---------------------------------------------------------------
 
-def run_agent(task: str, image_data_url: str | None = None, excluded: List[str] | None = None,
-              original_formats: List[str] | None = None, image_meta: str | None = None, 
-              conversation_history: List[str] | None = None) -> AgentToolSelection:
+def run_agent(
+    task: str,
+    image_data_url: str | None = None,
+    excluded: List[str] | None = None,
+    original_formats: List[str] | None = None,
+    image_meta: str | None = None,
+    conversation_history: List[str] | None = None,
+    model: str | None = None,
+    base_url: str | None = None,
+    top_k: int | None = None,
+    num_choices: int | None = None,
+) -> AgentToolSelection:
     """Execute the agent. We inline the image as extra context in user message (multimodal reasoning)."""
     extra_context = ""
     if image_data_url:
@@ -162,9 +175,15 @@ def run_agent(task: str, image_data_url: str | None = None, excluded: List[str]
 
     tool_logs: List[ToolRunLog] = []
 
-    # Intercept tool usage by patching agent? Simpler: rely on return types (pydantic-ai tracks internally, we record manually not available yet) -> for Phase 1 we skip deep logging.
-
-    deps = AgentState(excluded_tools=excluded or [])
+    # Create AgentState with runtime overrides
+    deps = AgentState(
+        excluded_tools=excluded or [],
+        override_model=model,
+        override_base_url=base_url,
+        override_top_k=top_k,
+        override_num_choices=num_choices,
+    )
+    
     # Provide hidden metadata context lines (non-user-visible) below a delimiter
     hidden_meta = ""
     if original_formats:
@@ -174,6 +193,10 @@ def run_agent(task: str, image_data_url: str | None = None, excluded: List[str]
         short_meta = " ".join(x.strip() for x in image_meta.splitlines() if x.strip())
         hidden_meta += "\n(Image Metadata: " + short_meta[:500] + ("…" if len(short_meta) > 500 else "") + ")"
 
+    # Add top_k hint if specified (for UI settings)
+    if top_k is not None:
+        hidden_meta += f"\n(Search top_k: {top_k})"
+    
     # Build prompt with conversation history if this is a follow-up
     if conversation_history and len(conversation_history) > 0:
         # Format previous conversation for context
@@ -182,11 +205,104 @@ def run_agent(task: str, image_data_url: str | None = None, excluded: List[str]
     else:
         prompt = task + extra_context + hidden_meta
 
-    result = agent.run_sync(prompt, deps=deps, output_type=ToolSelection, usage_limits=UsageLimits(tool_calls_limit=10)).output
+    # Determine which agent instance to use
+    agent_instance = agent  # Default to global agent
+    effective_num_choices = num_choices if num_choices is not None else 3
+    effective_model = model if model else agent_model_config.name
+    effective_top_k = top_k if top_k is not None else 12
+    
+    # When model is provided from UI, base_url comes with it (can be None for OpenAI)
+    # When model is NOT provided, use config defaults
+    if model:
+        # Model selected from dropdown - base_url parameter is authoritative
+        if base_url and "inference.rcp.epfl.ch" in base_url:
+            # EPFL model selected
+            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:
+            # OpenAI or other model selected (base_url=None means OpenAI)
+            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  # Will be 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
+        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"
+    log.info(
+        f"🤖 Agent execution - Model: {effective_model}, endpoint: {endpoint_display}, "
+        f"top_k: {effective_top_k}, num_choices: {effective_num_choices}, excluded: {len(excluded or [])}"
+    )
+    
+    # Create dynamic agent:
+    needs_dynamic_agent = (
+        (model and model != agent_model_config.name) or
+        (base_url is not None and base_url != agent_model_config.base_url) or
+        (runtime_api_key != api_key)  # API key mismatch - need new agent!
+    )
+    
+    if needs_dynamic_agent:
+        log.info(f"📦 Creating runtime agent with model={effective_model}, endpoint={effective_base_url or 'api.openai.com'}")
+        
+        runtime_provider = OpenAIProvider(
+            base_url=effective_base_url,
+            api_key=runtime_api_key,
+        )
+        runtime_model = OpenAIChatModel(model_name=effective_model, provider=runtime_provider)
+        agent_instance = Agent(
+            model=runtime_model,
+            system_prompt=get_agent_system_prompt(effective_num_choices),
+            deps_type=AgentState,
+        )
+        # Register tools on the dynamic agent
+        agent_instance.tool(search_tools, retries=2, prepare=cap_prepare)
+        agent_instance.tool(rerank, retries=2, prepare=cap_prepare)
+        agent_instance.tool(repo_info, retries=0, prepare=cap_prepare)
+        agent_instance.tool(resolve_demo_link, retries=2, prepare=cap_prepare)
+    elif num_choices is not None and num_choices != 3:
+        # Model/base_url same but num_choices differs - create agent with updated prompt
+        log.info(f"📦 Creating runtime agent with num_choices={effective_num_choices} (model: {effective_model})")
+        agent_instance = Agent(
+            model=openai_model,
+            system_prompt=get_agent_system_prompt(effective_num_choices),
+            deps_type=AgentState,
+        )
+        # Register tools on the dynamic agent
+        agent_instance.tool(search_tools, retries=2, prepare=cap_prepare)
+        agent_instance.tool(rerank, retries=2, prepare=cap_prepare)
+        agent_instance.tool(repo_info, retries=0, prepare=cap_prepare)
+        agent_instance.tool(resolve_demo_link, retries=2, prepare=cap_prepare)
+    else:
+        log.info(f"♻️  Using global agent (model: {effective_model}, num_choices: {effective_num_choices})")
+    
+    log.debug(f"Prompt length: {len(prompt)} chars, has_image: {image_data_url is not None}")
+    result = agent_instance.run_sync(prompt, deps=deps, output_type=ToolSelection, usage_limits=UsageLimits(tool_calls_limit=10)).output
+    log.info(f"✅ Agent execution complete - choices returned: {len(result.choices)}")
 
     # Convert tool call dicts into ToolRunLog entries
     for tc in deps.tool_calls:
-        tool_logs.append(ToolRunLog(tool=tc.get("tool"), inputs={k: v for k, v in tc.items() if k not in {"tool"}}, summary=str(tc)))
+        tool_logs.append(ToolRunLog(
+            tool=tc.get("tool"), 
+            inputs={k: v for k, v in tc.items() if k not in {"tool", "timestamp"}}, 
+            summary=str(tc),
+            timestamp=tc.get("timestamp")
+        ))
 
     # Post-run enrichment: pull demo links from resolve_demo_link tool calls
     demo_map = {}

src/ai_agent/agent/models.py

@@ -3,13 +3,14 @@
 from typing import List, Optional, Any, Dict
 from pydantic import BaseModel, Field
 
-from generator.schema import ToolChoice, ToolSelection, Conversation, ConversationStatus, CandidateDoc
+from ai_agent.generator.schema import ToolSelection, CandidateDoc
 
 class ToolRunLog(BaseModel):
     tool: str
     inputs: Dict[str, Any] = Field(default_factory=dict)
     summary: str
     error: Optional[str] = None
+    timestamp: Optional[str] = None
 
 class AgentToolSelection(ToolSelection):
     tool_calls: List[ToolRunLog] = Field(default_factory=list)

src/ai_agent/agent/tools/gradio_space_tool.py

@@ -4,8 +4,8 @@
 from pydantic import BaseModel
 import os, re, logging
 from .utils import get_pipeline
-from utils.utils import _best_runnable_link
-from utils.previews import _build_preview_for_vlm
+from ai_agent.utils.utils import _best_runnable_link
+from ai_agent.utils.previews import _build_preview_for_vlm
 from gradio_client import Client, handle_file
 import tempfile
 from pathlib import Path

src/ai_agent/agent/tools/rerank_tool.py

@@ -4,7 +4,7 @@
 from pydantic import BaseModel
 import os, re
 
-from retriever.software_doc import SoftwareDoc
+from ai_agent.retriever.software_doc import SoftwareDoc
 from .utils import get_pipeline
 
 class RerankInput(BaseModel):

src/ai_agent/agent/tools/search_tool.py

@@ -3,7 +3,7 @@
 from typing import List
 from pydantic import BaseModel, Field
 
-from generator.schema import CandidateDoc
+from ai_agent.generator.schema import CandidateDoc
 from .utils import get_pipeline
 
 class SearchToolsInput(BaseModel):