from __future__ import annotations import os import logging import time import asyncio import threading from collections import OrderedDict 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 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, Conversation, ConversationStatus from ai_agent.utils.config import get_config from .models import AgentToolSelection, ToolRunLog, UsageStats from .tools.repo_info_tool import tool_repo_summary, RepoSummaryInput from ai_agent.agent.utils import coerce_github_url_or_none from .tools.search_tool import tool_search_tools, SearchToolsInput from .tools.search_alternative_tool import ( tool_search_alternative, SearchAlternativeInput, ) from .tools.sparql_tool import tool_sparql_query, SparqlQueryInput from .tools.query_utils import sanitize_retrieval_query from .utils import AgentState, limit_tool_calls, cap_prepare from ai_agent.utils.image_meta import summarize_image_metadata, detect_ext_token log = logging.getLogger("agent.core") DEFAULT_NUM_CHOICES = int(os.getenv("NUM_CHOICES", "3")) # --------------------------------------------------------------------------- # Dynamic agent instance cache # Key: (model_name, base_url, api_key_env, num_choices) # Avoids rebuilding Agent/OpenAIProvider/model objects on every request when # the UI repeatedly uses the same custom endpoint + model combination. # Bounded LRU (max AGENT_CACHE_MAX entries); protected by a lock so that # concurrent requests cannot race while creating/inserting agents. # --------------------------------------------------------------------------- _AGENT_CACHE_MAX: int = int(os.getenv("AGENT_CACHE_MAX", "16")) _AGENT_CACHE_LOCK: threading.Lock = threading.Lock() _AGENT_CACHE: OrderedDict[tuple, "Agent"] = OrderedDict() # --------------------------------------------------------------------------- # Model / provider setup # --------------------------------------------------------------------------- config = get_config() agent_model_config = config.agent_model try: api_key = agent_model_config.get_api_key() except ValueError as e: log.error(f"Failed to get API key for agent model: {e}") raise log.info(f"Initializing agent model: {agent_model_config.name}") if agent_model_config.base_url: log.info(f"Using custom OpenAI base URL: {agent_model_config.base_url}") log.info("Using OpenAIChatModel (chat/completions API) for custom endpoint") provider = OpenAIProvider( 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) openai_model = OpenAIResponsesModel( model_name=agent_model_config.name, provider=provider, ) # --------------------------------------------------------------------------- # Agent definition # --------------------------------------------------------------------------- agent = Agent( model=openai_model, system_prompt=get_agent_system_prompt(DEFAULT_NUM_CHOICES), deps_type=AgentState, output_retries=int(os.getenv("AGENT_OUTPUT_RETRIES", "3")), ) # --------------------------------------------------------------------------- # Tool adapters for the agent # --------------------------------------------------------------------------- @agent.tool(retries=2, prepare=cap_prepare) @limit_tool_calls("search_tools", cap=1) async def search_tools( ctx: RunContext[AgentState], query: str, excluded: List[str] | None = None, top_k: int = 12, ) -> List[dict]: """ Agent-facing search tool. Delegates to tools.search_tool.tool_search_tools(), but automatically injects: - globally excluded tools (from ctx.deps.excluded_tools) - image_paths and original_formats (from ctx.deps, set in run_agent) so the language model never has to reason about file paths directly. """ # Merge explicit exclusions with global exclusions from AgentState explicit_excluded = excluded or [] global_excluded = getattr(ctx.deps, "excluded_tools", []) or [] all_excluded = sorted(set(explicit_excluded + list(global_excluded))) original_formats = getattr(ctx.deps, "original_formats", []) or [] image_paths = getattr(ctx.deps, "image_paths", []) or [] effective_top_k = ( ctx.deps.override_top_k if ctx.deps.override_top_k is not None else top_k ) started = time.perf_counter() inp = SearchToolsInput( query=sanitize_retrieval_query(query), excluded=all_excluded, top_k=effective_top_k, original_formats=original_formats, image_paths=image_paths, ) out = tool_search_tools(inp) ctx.deps.tool_calls.append( { "tool": "search_tools", "query": query, "count": len(out.candidates), "duration_ms": round((time.perf_counter() - started) * 1000, 1), "original_formats": original_formats, "excluded": all_excluded, "timestamp": datetime.now().isoformat(), } ) # Return plain dicts so the LLM sees a simple JSON-like structure. return [c.model_dump(mode="python") for c in out.candidates] @agent.tool(retries=2, prepare=cap_prepare) @limit_tool_calls("search_alternative", cap=3) async def search_alternative( ctx: RunContext[AgentState], alternative_query: str, excluded: List[str] | None = None, top_k: int = 12, ) -> List[dict]: """ Search with an alternative query formulation (includes automatic reranking). """ explicit_excluded = excluded or [] global_excluded = getattr(ctx.deps, "excluded_tools", []) or [] all_excluded = sorted(set(explicit_excluded + list(global_excluded))) original_formats = getattr(ctx.deps, "original_formats", []) or [] image_paths = getattr(ctx.deps, "image_paths", []) or [] started = time.perf_counter() inp = SearchAlternativeInput( alternative_query=alternative_query, excluded=all_excluded, top_k=top_k, original_formats=original_formats, image_paths=image_paths, ) out = tool_search_alternative(inp) ctx.deps.tool_calls.append( { "tool": "search_alternative", "alternative_query": alternative_query, "query_used": out.query_used, "count": len(out.candidates), "duration_ms": round((time.perf_counter() - started) * 1000, 1), "original_formats": original_formats, "excluded": all_excluded, "timestamp": datetime.now().isoformat(), } ) return [c.model_dump(mode="python") for c in out.candidates] @agent.tool(retries=1, prepare=cap_prepare) @limit_tool_calls("sparql_query", cap=3) async def sparql_query( ctx: RunContext[AgentState], query: str, limit: int = 50, ) -> dict: """Run a read-only SPARQL SELECT / ASK against the GraphDB catalog. Use this for questions the semantic search can't answer cleanly — aggregates ("how many tools support DICOM?"), structural filters ("tools that require a GPU AND are free"), distinct property values ("which licences appear in the catalog?"), etc. UPDATE-style operations are rejected. Results capped at `limit` rows (max 200). Returns columns + rows, plus a boolean for ASK queries. """ started = time.perf_counter() inp = SparqlQueryInput(query=query, limit=max(1, min(200, int(limit)))) out = tool_sparql_query(inp) ctx.deps.tool_calls.append( { "tool": "sparql_query", "query": query[:300], "row_count": out.row_count, "truncated": out.truncated, "boolean": out.boolean, "error": out.error, "duration_ms": round((time.perf_counter() - started) * 1000, 1), "timestamp": datetime.now().isoformat(), } ) return out.model_dump(mode="python") @agent.tool(retries=2, prepare=cap_prepare) @limit_tool_calls("repo_info_batch", cap=4) async def repo_info_batch( ctx: RunContext[AgentState], urls: List[str], ) -> List[dict]: """Fetch repository summaries for multiple repositories in parallel.""" started = time.perf_counter() if not urls: return [] normalized: List[str] = [] skipped: List[dict] = [] seen: set[str] = set() for raw in urls: norm = coerce_github_url_or_none(raw) if not norm: skipped.append( { "url": raw, "skipped": True, "reason": "NON_GITHUB_URL", } ) continue if norm in seen: continue seen.add(norm) normalized.append(norm) tasks = [tool_repo_summary(RepoSummaryInput(url=u)) for u in normalized] outcomes = await asyncio.gather(*tasks, return_exceptions=True) results: List[dict] = [] for url, outcome in zip(normalized, outcomes): if isinstance(outcome, Exception): results.append( { "url": url, "source": "error", "error": str(outcome), } ) continue payload = outcome.model_dump(mode="python") payload["url"] = url results.append(payload) if skipped: results.extend(skipped) ctx.deps.tool_calls.append( { "tool": "repo_info_batch", "requested": len(urls), "normalized": len(normalized), "returned": len(results), "duration_ms": round((time.perf_counter() - started) * 1000, 1), "timestamp": datetime.now().isoformat(), } ) return results # --------------------------------------------------------------------------- # High level entry point: run the agent on (text query + image) # --------------------------------------------------------------------------- def run_agent( task: str, image_paths: List[str], excluded: List[str] | None = None, conversation_history: List[str] | None = None, *, 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, ) -> AgentToolSelection: """ Execute the agent for a user task and at least one image path. - derive canonical original_formats (tiff / dicom / nifti / ...) - build a compact image metadata summary (or use pre-computed one) - pass both to the LLM as hidden context - store image_paths/original_formats in deps so retrieval tools can use them - optionally allow runtime model/base_url/top_k/num_choices overrides IMPORTANT: The model only sees an actual image if `image_bytes` is provided. `image_paths` are used for metadata + tool context only. """ run_started = time.perf_counter() # image_paths may now be empty — the agent runs in text-only mode and # skips the VLM. Retrieval still works on the text query alone. image_paths = list(image_paths or []) tool_logs: List[ToolRunLog] = [] # ---- 1) Derive image-based metadata and format hints -------------------- metadata_started = time.perf_counter() meta_str = ( image_metadata if image_metadata is not None else (summarize_image_metadata(image_paths) or "") ) fmt_str = detect_ext_token(image_paths) or "" original_formats = [t.lower() for t in fmt_str.split()] if fmt_str else [] metadata_duration_ms = round((time.perf_counter() - metadata_started) * 1000, 1) effective_top_k = top_k if top_k is not None else 12 effective_num_choices = num_choices if num_choices is not None else 3 # ---- 2) Prepare dependency state passed to all tools -------------------- deps = AgentState( excluded_tools=excluded or [], override_model=model, override_base_url=base_url, override_top_k=effective_top_k, override_num_choices=effective_num_choices, ) setattr(deps, "image_paths", list(image_paths)) setattr(deps, "original_formats", original_formats) # ---- 3) Hidden metadata lines for the model ---------------------------- hidden_meta = "" if original_formats: hidden_meta += "\n(Formats Hint: " + ",".join(original_formats) + ")" if meta_str: short_meta = " ".join(x.strip() for x in meta_str.splitlines() if x.strip()) hidden_meta += ( "\n(Image Metadata: " + short_meta[:500] + ("…" if len(short_meta) > 500 else "") + ")" ) hidden_meta += f"\n(Search top_k: {effective_top_k})" extra_context = "\n\n**CRITICAL: Analyze the attached preview image showing the user's data.**\nUse visual observations (anatomy visible, image quality, dimensionality, contrast) combined with the metadata below to recommend tools. Reference what you see in your explanations." # ---- 4) Build the prompt (optionally including history) ---------------- if conversation_history and len(conversation_history) > 0: history_text = "\n".join(conversation_history) prompt = ( f"Previous conversation:\n{history_text}\n\n" f"Current request: {task}{extra_context}{hidden_meta}" ) else: prompt = task + extra_context + hidden_meta # ----------------------------------------------------------------------- # Determine which agent instance to use # ----------------------------------------------------------------------- agent_instance = 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) 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} is set: {bool(runtime_api_key)}") 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}") # 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 [])}" ) needs_dynamic_agent = model is not None if needs_dynamic_agent: cache_key = (effective_model, effective_base_url or "", api_key_env or "OPENAI_API_KEY", effective_num_choices) with _AGENT_CACHE_LOCK: agent_instance = _AGENT_CACHE.get(cache_key) if agent_instance is not None: _AGENT_CACHE.move_to_end(cache_key) if agent_instance is None: 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, ) # Use OpenAIChatModel (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, system_prompt=get_agent_system_prompt(effective_num_choices), deps_type=AgentState, output_retries=int(os.getenv("AGENT_OUTPUT_RETRIES", "3")), ) # Register tools on the dynamic agent agent_instance.tool(search_tools, retries=2, prepare=cap_prepare) agent_instance.tool(search_alternative, retries=2, prepare=cap_prepare) agent_instance.tool(repo_info_batch, retries=2, prepare=cap_prepare) agent_instance.tool(sparql_query, retries=1, prepare=cap_prepare) with _AGENT_CACHE_LOCK: _AGENT_CACHE[cache_key] = agent_instance _AGENT_CACHE.move_to_end(cache_key) while len(_AGENT_CACHE) > _AGENT_CACHE_MAX: _AGENT_CACHE.popitem(last=False) else: log.info( f"♻️ Reusing cached dynamic agent (model: {effective_model}, num_choices: {effective_num_choices})" ) elif ( num_choices is not None and num_choices != DEFAULT_NUM_CHOICES ): cache_key = (effective_model, effective_base_url or "", api_key_env or "OPENAI_API_KEY", effective_num_choices) with _AGENT_CACHE_LOCK: agent_instance = _AGENT_CACHE.get(cache_key) if agent_instance is not None: _AGENT_CACHE.move_to_end(cache_key) if agent_instance is None: 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, output_retries=int(os.getenv("AGENT_OUTPUT_RETRIES", "3")), ) # Register tools on the dynamic agent agent_instance.tool(search_tools, retries=2, prepare=cap_prepare) agent_instance.tool(search_alternative, retries=2, prepare=cap_prepare) agent_instance.tool(repo_info_batch, retries=2, prepare=cap_prepare) agent_instance.tool(sparql_query, retries=1, prepare=cap_prepare) with _AGENT_CACHE_LOCK: _AGENT_CACHE[cache_key] = agent_instance _AGENT_CACHE.move_to_end(cache_key) while len(_AGENT_CACHE) > _AGENT_CACHE_MAX: _AGENT_CACHE.popitem(last=False) else: log.info( f"♻️ Reusing cached dynamic agent with num_choices={effective_num_choices} (model: {effective_model})" ) 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_paths: {bool(image_paths)}, has_image_bytes: {bool(image_bytes)}" ) # ---- 5) Build multimodal prompt if image bytes provided ---------------- if image_bytes: log.info( f"🖼️ Sending image preview to model ({len(image_bytes)} bytes = {len(image_bytes)/1024:.1f}KB)" ) user_prompt = [ prompt, BinaryContent( data=image_bytes, media_type="image/png", ), ] else: log.warning( "⚠️ No image bytes provided - the model will not see the image preview" ) user_prompt = prompt # ---- 6) Run the agent -------------------------------------------------- try: llm_started = time.perf_counter() run_result = agent_instance.run_sync( user_prompt, deps=deps, output_type=ToolSelection, usage_limits=UsageLimits(tool_calls_limit=20), ) llm_duration_ms = round((time.perf_counter() - llm_started) * 1000, 1) result = run_result.output 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"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." ) except Exception as e: # Handle global tool quota limit (UsageLimitExceeded) and other errors gracefully error_msg = str(e) llm_duration_ms = round((time.perf_counter() - llm_started) * 1000, 1) log.warning(f"⚠️ Agent execution encountered an error: {error_msg}") run_result = None # Ensure run_result is defined for usage stats extraction # 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", []): tool_name = tc.get("tool") timestamp = tc.get("timestamp") error = tc.get("error") inputs = { k: v for k, v in tc.items() if k not in ("tool", "timestamp", "error") } tool_logs.append( ToolRunLog( tool=tool_name, inputs=inputs, timestamp=timestamp, error=error, ) ) stage_counts: dict[str, int] = {} stage_durations: dict[str, float] = {} for tc in getattr(deps, "tool_calls", []): name = tc.get("tool", "unknown") stage_counts[name] = stage_counts.get(name, 0) + 1 duration_ms = tc.get("duration_ms") if isinstance(duration_ms, (int, float)): stage_durations[name] = stage_durations.get(name, 0.0) + float(duration_ms) total_duration_ms = round((time.perf_counter() - run_started) * 1000, 1) log.info( "⏱️ Latency summary: total_ms=%s metadata_ms=%s llm_ms=%s tools=%s tool_ms=%s", total_duration_ms, metadata_duration_ms, llm_duration_ms, stage_counts, {k: round(v, 1) for k, v in stage_durations.items()}, ) # ---- 8) Extract usage statistics if available ------------------------- usage_stats = None if run_result and hasattr(run_result, "usage") and run_result.usage: usage = run_result.usage() usage_stats = UsageStats( total_tokens=usage.total_tokens, input_tokens=usage.input_tokens, output_tokens=usage.output_tokens, ) # ---- 9) Wrap into high-level AgentToolSelection ------------------------ return AgentToolSelection( conversation=result.conversation, choices=result.choices, explanation=result.explanation, reason=result.reason, tool_calls=tool_logs, usage=usage_stats, ) __all__ = ["run_agent", "agent"]