code-puppy 0.0.355__py3-none-any.whl → 0.0.357__py3-none-any.whl

code_puppy/tools/browser/vqa_agent.py

@@ -2,10 +2,12 @@
 
 from __future__ import annotations
 
-from functools import lru_cache
+from collections.abc import AsyncIterable
+from typing import Any
 
 from pydantic import BaseModel, Field
-from pydantic_ai import Agent, BinaryContent
+from pydantic_ai import Agent, BinaryContent, PartDeltaEvent, PartStartEvent, RunContext
+from pydantic_ai.messages import TextPart, TextPartDelta
 
 from code_puppy.config import get_use_dbos, get_vqa_model_name
 
@@ -18,73 +20,175 @@ class VisualAnalysisResult(BaseModel):
     observations: str
 
 
-def _get_vqa_instructions() -> str:
-    """Get the system instructions for the VQA agent."""
-    return (
-        "You are a visual analysis specialist. Answer the user's question about the provided image. "
-        "Always respond using the structured schema: answer, confidence (0-1 float), observations. "
-        "Confidence reflects how certain you are about the answer. Observations should include useful, concise context."
-    )
+DEFAULT_VQA_INSTRUCTIONS = (
+    "You are a visual analysis specialist. Answer the user's question about the provided image. "
+    "Always respond using the structured schema: answer, confidence (0-1 float), observations. "
+    "Confidence reflects how certain you are about the answer. Observations should include useful, concise context."
+)
 
 
-@lru_cache(maxsize=1)
-def _load_vqa_agent(model_name: str) -> Agent[None, VisualAnalysisResult]:
-    """Create a cached agent instance for visual analysis."""
+async def run_vqa_analysis(
+    question: str,
+    image_bytes: bytes,
+    media_type: str = "image/png",
+) -> str:
+    """Execute the VQA agent asynchronously against screenshot bytes.
+
+    Follows the same pattern as agent_tools.py for prompt preparation
+    and model configuration.
+
+    Args:
+        question: The question to ask about the image.
+        image_bytes: The raw image bytes.
+        media_type: The MIME type of the image (default: "image/png").
+        system_prompt: Optional custom system prompt. If None, uses default VQA instructions.
+
+    Returns:
+        str: The answer from the VQA analysis.
+    """
+    from code_puppy import callbacks
     from code_puppy.model_factory import ModelFactory
     from code_puppy.model_utils import prepare_prompt_for_model
 
+    # Get model configuration
+    model_name = get_vqa_model_name()
     models_config = ModelFactory.load_config()
     model = ModelFactory.get_model(model_name, models_config)
 
-    # Handle claude-code models: swap instructions (prompt prepending happens in run_vqa_analysis)
-    instructions = _get_vqa_instructions()
+    # Build instructions: custom system_prompt or default VQA instructions
+    instructions = DEFAULT_VQA_INSTRUCTIONS
+
+    # Apply prompt additions (like file permission handling) - same as agent_tools.py
+    prompt_additions = callbacks.on_load_prompt()
+    if prompt_additions:
+        instructions += "\n" + "\n".join(prompt_additions)
+
+    # Handle claude-code models: swap instructions, prepend system prompt to user question
+    # Following the exact pattern from agent_tools.py
     prepared = prepare_prompt_for_model(
-        model_name, instructions, "", prepend_system_to_user=False
+        model_name, instructions, question, prepend_system_to_user=True
     )
     instructions = prepared.instructions
+    question = prepared.user_prompt
 
+    # Create the VQA agent with string output
     vqa_agent = Agent(
         model=model,
         instructions=instructions,
-        output_type=VisualAnalysisResult,
-        retries=2,
     )
 
+    # Wrap with DBOS if enabled
     if get_use_dbos():
         from pydantic_ai.durable_exec.dbos import DBOSAgent
 
-        dbos_agent = DBOSAgent(vqa_agent, name="vqa-agent")
-        return dbos_agent
+        vqa_agent = DBOSAgent(vqa_agent, name="vqa-agent")
 
-    return vqa_agent
+    # Run the agent with the image
+    result = await vqa_agent.run(
+        [
+            question,
+            BinaryContent(data=image_bytes, media_type=media_type),
+        ]
+    )
+    return result.output
 
 
-def _get_vqa_agent() -> Agent[None, VisualAnalysisResult]:
-    """Return a cached VQA agent configured with the current model."""
-    model_name = get_vqa_model_name()
-    # lru_cache keyed by model_name ensures refresh when configuration changes
-    return _load_vqa_agent(model_name)
+def _create_vqa_stream_handler(
+    accumulator: list[str],
+):
+    """Create an event stream handler that accumulates text.
+
+    Args:
+        accumulator: List to accumulate text chunks into (pass empty list).
+
+    Returns:
+        Async event stream handler function.
+    """
+
+    async def vqa_event_stream_handler(
+        ctx: RunContext,
+        events: AsyncIterable[Any],
+    ) -> None:
+        """Handle streaming events - print text as it arrives."""
+        async for event in events:
+            # Handle text part start - might have initial content
+            if isinstance(event, PartStartEvent):
+                if isinstance(event.part, TextPart) and event.part.content:
+                    accumulator.append(event.part.content)
 
+            # Handle text deltas - the streaming bits
+            elif isinstance(event, PartDeltaEvent):
+                if isinstance(event.delta, TextPartDelta) and event.delta.content_delta:
+                    accumulator.append(event.delta.content_delta)
 
-def run_vqa_analysis(
+    return vqa_event_stream_handler
+
+
+async def run_vqa_analysis_stream(
     question: str,
     image_bytes: bytes,
     media_type: str = "image/png",
-) -> VisualAnalysisResult:
-    """Execute the VQA agent synchronously against screenshot bytes."""
-    from code_puppy.model_utils import prepare_prompt_for_model
+) -> str:
+    """Execute the VQA agent with streaming output.
 
-    agent = _get_vqa_agent()
+    Streams text to console as it arrives and accumulates the full response.
 
-    # Handle claude-code models: prepend system prompt to user question
+    Args:
+        question: The question to ask about the image.
+        image_bytes: The raw image bytes.
+        media_type: The MIME type of the image (default: "image/png").
+
+    Returns:
+        str: The accumulated answer from the VQA analysis.
+    """
+    from code_puppy import callbacks
+    from code_puppy.model_factory import ModelFactory
+    from code_puppy.model_utils import prepare_prompt_for_model
+
+    # Get model configuration
     model_name = get_vqa_model_name()
-    prepared = prepare_prompt_for_model(model_name, _get_vqa_instructions(), question)
+    models_config = ModelFactory.load_config()
+    model = ModelFactory.get_model(model_name, models_config)
+
+    # Build instructions
+    instructions = DEFAULT_VQA_INSTRUCTIONS
+
+    # Apply prompt additions (like file permission handling)
+    prompt_additions = callbacks.on_load_prompt()
+    if prompt_additions:
+        instructions += "\n" + "\n".join(prompt_additions)
+
+    # Handle claude-code models: swap instructions, prepend system prompt to user question
+    prepared = prepare_prompt_for_model(
+        model_name, instructions, question, prepend_system_to_user=True
+    )
+    instructions = prepared.instructions
     question = prepared.user_prompt
 
-    result = agent.run_sync(
+    # Create the VQA agent
+    vqa_agent = Agent(
+        model=model,
+        instructions=instructions,
+    )
+
+    # Wrap with DBOS if enabled
+    if get_use_dbos():
+        from pydantic_ai.durable_exec.dbos import DBOSAgent
+
+        vqa_agent = DBOSAgent(vqa_agent, name="vqa-agent-stream")
+
+    # Accumulator for streamed text (use list to allow mutation in handler)
+    accumulated_chunks: list[str] = []
+
+    # Create the stream handler
+    stream_handler = _create_vqa_stream_handler(accumulated_chunks)
+
+    # Run the agent with event_stream_handler
+    result = await vqa_agent.run(
         [
             question,
             BinaryContent(data=image_bytes, media_type=media_type),
-        ]
+        ],
+        event_stream_handler=stream_handler,
     )
     return result.output