entari-plugin-hyw 4.0.0rc5__py3-none-any.whl

entari_plugin_hyw/stage_instruct.py

@@ -0,0 +1,328 @@
+"""
+Instruct Stage
+
+Handles initial task planning and search generation.
+Analyze user query and execute initial searches.
+"""
+
+import json
+import time
+import asyncio
+from typing import Any, Dict, List, Optional, Tuple
+from loguru import logger
+from openai import AsyncOpenAI
+
+from .stage_base import BaseStage, StageContext, StageResult
+from .definitions import (
+    get_refuse_answer_tool,
+    get_web_search_tool,
+    get_crawl_page_tool,
+    INSTRUCT_SP
+)
+
+class InstructStage(BaseStage):
+    @property
+    def name(self) -> str:
+        return "Instruct"
+    
+    def __init__(self, config: Any, search_service: Any, client: AsyncOpenAI):
+        super().__init__(config, search_service, client)
+        
+        self.refuse_answer_tool = get_refuse_answer_tool()
+        self.web_search_tool = get_web_search_tool()
+        self.crawl_page_tool = get_crawl_page_tool()
+    
+    async def execute(self, context: StageContext) -> StageResult:
+        start_time = time.time()
+        
+        # --- Round 1: Initial Discovery ---
+        logger.info("Instruct: Starting Round 1 (Initial Discovery)")
+        
+        # Build Round 1 User Message
+        r1_user_content = self._build_user_message(context)
+        r1_messages = [
+            {"role": "system", "content": INSTRUCT_SP},
+            {"role": "user", "content": r1_user_content}
+        ]
+        
+        # Execute Round 1 LLM
+        r1_response, r1_usage, r1_tool_calls, r1_content = await self._call_llm(
+            messages=r1_messages,
+            tools=[self.refuse_answer_tool, self.web_search_tool, self.crawl_page_tool],
+            tool_choice="auto"
+        )
+        
+        if context.should_refuse:
+             # If refused in Round 1, stop here
+             return self._build_result(start_time, r1_usage, r1_content, len(r1_tool_calls or []))
+
+        # Execute Round 1 Tools
+        r1_tool_outputs = []
+        if r1_tool_calls:
+            r1_tool_outputs = await self._process_tool_calls(context, r1_tool_calls)
+        
+        # --- Context Assembly for Round 2 ---
+        
+        # Summarize Round 1 actions for context
+        r1_summary_text = "## Round 1 Execution Summary\n"
+        if r1_content:
+            r1_summary_text += f"Thought: {r1_content}\n"
+            
+        if r1_tool_outputs:
+            r1_summary_text += "Tools Executed & Results:\n"
+            for output in r1_tool_outputs:
+                # content here is the tool output (e.g. search results text or crawl preview)
+                r1_summary_text += f"- Action: {output['name']}\n"
+                r1_summary_text += f"  Result: {output['content']}\n"
+        else:
+            r1_summary_text += "No tools were executed in Round 1.\n"
+            
+        r2_context_str = f"""User Query: {context.user_input}
+        
+{r1_summary_text}
+"""
+        # Save to context for next stage
+        context.review_context = r2_context_str
+        
+        # Update instruct_history for logging/record purposes
+        context.instruct_history.append({
+            "role": "assistant",
+            "content": f"[Round 1 Thought]: {r1_content}\n[Round 1 Actions]: {len(r1_tool_outputs)} tools"
+        })
+        
+        return self._build_result(start_time, r1_usage, r1_content, len(r1_tool_calls or []))
+
+    def _build_user_message(self, context: StageContext) -> Any:
+        text_prompt = f"User Query: {context.user_input}"
+        if context.images:
+            user_content: List[Dict[str, Any]] = [{"type": "text", "text": text_prompt}]
+            for img_b64 in context.images:
+                url = f"data:image/jpeg;base64,{img_b64}" if not img_b64.startswith("data:") else img_b64
+                user_content.append({"type": "image_url", "image_url": {"url": url}})
+            return user_content
+        return text_prompt
+
+    async def _call_llm(self, messages, tools, tool_choice="auto"):
+        model_cfg = self.config.get_model_config("instruct")
+        client = self._client_for(
+            api_key=model_cfg.get("api_key"),
+            base_url=model_cfg.get("base_url")
+        )
+        model = model_cfg.get("model_name") or self.config.model_name
+        
+        try:
+            response = await client.chat.completions.create(
+                model=model,
+                messages=messages,
+                tools=tools,
+                tool_choice=tool_choice,
+                temperature=self.config.temperature,
+                extra_body=model_cfg.get("extra_body"),
+            )
+        except Exception as e:
+            logger.error(f"InstructStage LLM Error: {e}")
+            raise e
+            
+        usage = {"input_tokens": 0, "output_tokens": 0}
+        if hasattr(response, "usage") and response.usage:
+            usage["input_tokens"] = getattr(response.usage, "prompt_tokens", 0) or 0
+            usage["output_tokens"] = getattr(response.usage, "completion_tokens", 0) or 0
+            
+        message = response.choices[0].message
+        content = message.content or ""
+        tool_calls = message.tool_calls
+        
+        if content:
+             logger.debug(f"Instruct: Agent Thought -> {content[:100]}...")
+             
+        return response, usage, tool_calls, content
+
+    async def _process_tool_calls(self, context: StageContext, tool_calls: List[Any]) -> List[Dict[str, Any]]:
+        """
+        Executes tool calls and returns a list of outputs for context building.
+        Updates context.web_results globally.
+        """
+        pending_crawls = []    # List of (url, tool_call_id)
+        pending_searches = []  # List of (query, tool_call_id)
+        
+        results_for_context = []
+
+        for tc in tool_calls:
+            name = tc.function.name
+            tc_id = tc.id
+            try:
+                args = json.loads(tc.function.arguments)
+            except json.JSONDecodeError:
+                results_for_context.append({
+                    "id": tc_id, "name": name, "content": "Error: Invalid JSON arguments"
+                })
+                continue
+            
+            if name == "refuse_answer":
+                reason = args.get("reason", "Refused")
+                logger.warning(f"Instruct: Model Refused Answer. Reason: {reason}")
+                context.should_refuse = True
+                context.refuse_reason = reason
+                results_for_context.append({
+                     "id": tc_id, "name": name, "content": f"Refused: {reason}"
+                })
+            
+            elif name == "internal_web_search":
+                query = args.get("query")
+                if query: 
+                    logger.info(f"Instruct: Planned search query -> '{query}'")
+                    pending_searches.append((query, tc_id))
+            
+            elif name == "crawl_page":
+                url = args.get("url")
+                if url: 
+                    logger.info(f"Instruct: Planned page crawl -> {url}")
+                    pending_crawls.append((url, tc_id))
+        
+        # Execute Batches
+        
+        # 1. Crawls
+        if pending_crawls:
+            urls = [u for u, _ in pending_crawls]
+            logger.info(f"Instruct: Executing {len(urls)} crawls via batch...")
+            
+            # Start fetch
+            fetch_task = asyncio.create_task(self.search_service.fetch_pages_batch(urls))
+            
+            is_image_mode = getattr(self.config, "page_content_mode", "text") == "image"
+            tab_ids = []
+            if is_image_mode:
+                from .render_vue import get_content_renderer
+                renderer = await get_content_renderer()
+                loop = asyncio.get_running_loop()
+                tab_tasks = [
+                    loop.run_in_executor(renderer._executor, renderer._prepare_tab_sync)
+                    for _ in urls
+                ]
+                tab_ids = await asyncio.gather(*tab_tasks, return_exceptions=True)
+                logger.debug(f"Instruct: Prepared {len(tab_ids)} tabs: {tab_ids}")
+            
+            crawl_results_list = await fetch_task
+            
+            if is_image_mode and tab_ids:
+                theme_color = getattr(self.config, "theme_color", "#ef4444")
+                render_tasks = []
+                valid_pairs = []
+                MAX_CHARS = 3000
+                for i, (page_data, tab_id) in enumerate(zip(crawl_results_list, tab_ids)):
+                    if isinstance(tab_id, Exception):
+                        logger.warning(f"Instruct: Skip rendering page {i} due to tab error: {tab_id}")
+                        continue
+                    
+                    # Truncate content to avoid excessive size
+                    content = page_data.get("content", "")
+                    if len(content) > MAX_CHARS:
+                        content = content[:MAX_CHARS] + "\n\n...(content truncated for length)..."
+                        page_data["content"] = content
+
+                    if not content:
+                        logger.warning(f"Instruct: Skip rendering page {i} due to empty content")
+                        continue
+                    
+                    valid_pairs.append((i, page_data))
+                    render_tasks.append(
+                        loop.run_in_executor(
+                            renderer._executor,
+                            renderer._render_page_to_b64_sync,
+                            {"title": page_data.get("title", "Page"), "content": content},
+                            tab_id,
+                            theme_color
+                        )
+                    )
+                
+                if render_tasks:
+                    logger.debug(f"Instruct: Parallel rendering {len(render_tasks)} pages...")
+                    screenshots = await asyncio.gather(*render_tasks, return_exceptions=True)
+                    logger.debug(f"Instruct: Parallel rendering finished. Results count: {len(screenshots)}")
+                    for j, (orig_idx, page_data) in enumerate(valid_pairs):
+                         if j < len(screenshots) and not isinstance(screenshots[j], Exception):
+                             crawl_results_list[orig_idx]["screenshot_b64"] = screenshots[j]
+
+            for i, (url, tc_id) in enumerate(pending_crawls):
+                page_data = crawl_results_list[i]
+                title = page_data.get("title", "Unknown")
+                
+                # Update global context
+                page_item = {
+                    "_id": context.next_id(),
+                    "_type": "page",
+                    "title": page_data.get("title", "Page"),
+                    "url": page_data.get("url", url),
+                    "content": page_data.get("content", ""),
+                    "is_crawled": True,
+                }
+                if page_data.get("screenshot_b64"):
+                    page_item["screenshot_b64"] = page_data["screenshot_b64"]
+                if page_data.get("raw_screenshot_b64"):
+                    page_item["raw_screenshot_b64"] = page_data["raw_screenshot_b64"]
+                if page_data.get("images"):
+                    page_item["images"] = page_data["images"]
+                    
+                context.web_results.append(page_item)
+                
+                # Output for Context Assembly
+                content_preview = page_data.get("content", "")[:500]
+                results_for_context.append({
+                     "id": tc_id, 
+                     "name": "crawl_page", 
+                     "content": f"Crawled '{title}' ({url}):\n{content_preview}..."
+                })
+
+        # 2. Searches
+        if pending_searches:
+            queries = [q for q, _ in pending_searches]
+            logger.info(f"Instruct: Executing {len(queries)} searches via batch...")
+            
+            search_results_list = await self.search_service.search_batch(queries)
+            
+            for i, (query, tc_id) in enumerate(pending_searches):
+                web_results = search_results_list[i]
+                visible_results = [r for r in web_results if not r.get("_hidden")]
+                
+                # Update global context
+                for item in web_results:
+                    item["_id"] = context.next_id()
+                    if "type" in item:
+                        item["_type"] = item["type"]
+                    elif "_type" not in item:
+                        item["_type"] = "search"
+                    item["query"] = query
+                    context.web_results.append(item)
+                
+                # Output for Context Assembly
+                summary = f"Found {len(visible_results)} results for '{query}':\n"
+                for r in visible_results[:5]:
+                    summary += f"- {r.get('title')} ({r.get('url')}): {(r.get('content') or '')[:100]}...\n"
+                
+                results_for_context.append({
+                     "id": tc_id, 
+                     "name": "internal_web_search", 
+                     "content": summary
+                })
+                
+        return results_for_context
+
+    def _build_result(self, start_time, usage, content, tool_calls_count):
+        model_cfg = self.config.get_model_config("instruct")
+        model = model_cfg.get("model_name") or self.config.model_name
+        
+        trace = {
+            "stage": "Instruct",
+            "model": model, 
+            "usage": usage,
+            "output": content,
+            "tool_calls": tool_calls_count,
+            "time": time.time() - start_time,
+        }
+        
+        return StageResult(
+            success=True,
+            data={"reasoning": content},
+            usage=usage,
+            trace=trace
+        )

entari_plugin_hyw/stage_instruct_review.py

@@ -0,0 +1,92 @@
+"""
+Instruct Review Stage
+
+Handles the second round of instruction: Review and Refine.
+Inherits from InstructStage to reuse tool execution logic.
+"""
+
+import time
+from typing import Any, List
+from loguru import logger
+from openai import AsyncOpenAI
+
+from .stage_base import StageContext, StageResult
+from .stage_instruct import InstructStage
+from .definitions import INSTRUCT_REVIEW_SP
+
+class InstructReviewStage(InstructStage):
+    @property
+    def name(self) -> str:
+        return "Instruct Review"
+    
+    def __init__(self, config: Any, search_service: Any, client: AsyncOpenAI):
+        super().__init__(config, search_service, client)
+        # Inherits tools from InstructStage
+    
+    async def execute(self, context: StageContext) -> StageResult:
+        start_time = time.time()
+        logger.info("Instruct Review: Starting Round 2 (Review & Refine)")
+        
+        # Check if we have context to review
+        if not context.review_context:
+            logger.warning("Instruct Review: No context found from Round 1. Skipping.")
+            return StageResult(success=True, data={"reasoning": "Skipped due to missing context."})
+            
+        # Build System Prompt (Clean)
+        system_prompt = INSTRUCT_REVIEW_SP
+        
+        # Build Messages
+        # Inject context as a separate user message explaining the background
+        context_message = f"## Previous Round Context\n\n```context\n{context.review_context}\n```"
+        
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": context_message},
+            {"role": "user", "content": self._build_user_message(context)}
+        ]
+        
+        # Call LLM
+        # We reuse _call_llm from parent
+        # We reuse tools from parent (refuse_answer might be redundant but harmless, or we can filter)
+        tools = [self.web_search_tool, self.crawl_page_tool] # Review prompt doesn't mention refuse_answer explicitly, but usually fine.
+        
+        response, usage, tool_calls, content = await self._call_llm(
+            messages=messages,
+            tools=tools,
+            tool_choice="auto"
+        )
+        
+        # Execute Tools
+        tool_outputs = []
+        if tool_calls:
+            tool_outputs = await self._process_tool_calls(context, tool_calls)
+            
+        # Update history logic? 
+        # The prompt says "上下文". It is "independent". 
+        # But for the record, we might want to log it.
+        context.instruct_history.append({
+             "role": "assistant",
+             "content": f"[Round 2 Review]: {content}\n[Round 2 Actions]: {len(tool_outputs)} tools"
+        })
+        
+        return self._build_result(start_time, usage, content, len(tool_calls or []))
+
+    def _build_result(self, start_time, usage, content, tool_calls_count):
+        model_cfg = self.config.get_model_config("instruct")
+        model = model_cfg.get("model_name") or self.config.model_name
+
+        trace = {
+            "stage": "Instruct Review",
+            "model": model, 
+            "usage": usage,
+            "output": content,
+            "tool_calls": tool_calls_count,
+            "time": time.time() - start_time,
+        }
+        
+        return StageResult(
+            success=True,
+            data={"reasoning": content},
+            usage=usage,
+            trace=trace
+        )

entari_plugin_hyw/stage_summary.py

@@ -0,0 +1,164 @@
+"""
+Summary Stage
+
+Generates final response based on gathered information.
+Different output formats for different modes.
+"""
+
+import time
+import re
+from typing import Any, Dict, List, Optional
+
+from loguru import logger
+from openai import AsyncOpenAI
+
+from .stage_base import BaseStage, StageContext, StageResult
+from .definitions import SUMMARY_REPORT_SP
+
+
+class SummaryStage(BaseStage):
+    """
+    Summary Stage: Generate final response.
+    """
+    
+    @property
+    def name(self) -> str:
+        return "Summary"
+    
+    async def execute(
+        self, 
+        context: StageContext, 
+        images: List[str] = None
+    ) -> StageResult:
+        """Generate summary."""
+        start_time = time.time()
+        
+        # Format context from web results
+        web_content = self._format_web_content(context)
+        full_context = f"{context.agent_context}\n\n{web_content}"
+        
+        # Select prompt
+        language = getattr(self.config, "language", "Simplified Chinese")
+        
+        system_prompt = SUMMARY_REPORT_SP.format(
+            language=language
+        )
+        
+        # Build Context Message
+        context_message = f"## Web Search & Page Content\n\n```context\n{full_context}\n```"
+        
+        # Build user content
+        user_text = context.user_input or "..."
+        if images:
+            user_content: List[Dict[str, Any]] = [{"type": "text", "text": user_text}]
+            for img_b64 in images:
+                url = f"data:image/jpeg;base64,{img_b64}" if not img_b64.startswith("data:") else img_b64
+                user_content.append({"type": "image_url", "image_url": {"url": url}})
+        else:
+            user_content = user_text
+        
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": context_message},
+            {"role": "user", "content": user_content}
+        ]
+        
+        # Get model config
+        model_cfg = self.config.get_model_config("main")
+        
+        client = self._client_for(
+            api_key=model_cfg.get("api_key"),
+            base_url=model_cfg.get("base_url")
+        )
+        
+        model = model_cfg.get("model_name") or self.config.model_name
+        
+        try:
+            response = await client.chat.completions.create(
+                model=model,
+                messages=messages,
+                temperature=self.config.temperature,
+                extra_body=getattr(self.config, "summary_extra_body", None),
+            )
+        except Exception as e:
+            logger.error(f"SummaryStage LLM error: {e}")
+            return StageResult(
+                success=False,
+                error=str(e),
+                data={"content": f"Error generating summary: {e}"}
+            )
+        
+        usage = {"input_tokens": 0, "output_tokens": 0}
+        if hasattr(response, "usage") and response.usage:
+            usage["input_tokens"] = getattr(response.usage, "prompt_tokens", 0) or 0
+            usage["output_tokens"] = getattr(response.usage, "completion_tokens", 0) or 0
+        
+        content = (response.choices[0].message.content or "").strip()
+        
+        return StageResult(
+            success=True,
+            data={"content": content},
+            usage=usage,
+            trace={
+                "model": model,
+                "provider": model_cfg.get("model_provider") or "Unknown",
+                "usage": usage,
+                "system_prompt": system_prompt,
+                "output": content,
+                "time": time.time() - start_time,
+                "images_count": len(images) if images else 0,
+            }
+        )
+    
+    def _strip_links(self, text: str) -> str:
+        """Strip markdown links [text](url) -> text and remove bare URLs."""
+        # Replace [text](url) with text
+        text = re.sub(r'\[([^\]]+)\]\([^\)]+\)', r'\1', text)
+        # Remove bare URLs (http/https) roughly, trying to preserve surrounding text if possible?
+        # A simple pattern for http/s
+        text = re.sub(r'https?://\S+', '', text)
+        return text
+
+    def _format_web_content(self, context: StageContext) -> str:
+        """Format web results for summary prompt."""
+        if not context.web_results:
+            return ""
+        
+        # Sort results: pages first, then raw searches, then snippets
+        def get_priority(item_type):
+            if item_type == "page": return 0
+            if item_type == "search_raw_page": return 1
+            return 2  # search (snippets)
+            
+        sorted_results = sorted(
+            context.web_results, 
+            key=lambda x: get_priority(x.get("_type"))
+        )
+        
+        lines = []
+        seen_urls = set()
+        
+        for res in sorted_results:
+            type_ = res.get("_type")
+            idx = res.get("_id")
+            title = (res.get("title", "") or "").strip()
+            url = res.get("url", "")
+            
+            # Deduplicate items by URL (keep higher priority item only)
+            if url:
+                if url in seen_urls:
+                    continue
+                seen_urls.add(url)
+            
+            # url = res.get("url", "") # Removed as requested
+            
+            if type_ == "page":
+                content = (res.get("content", "") or "").strip()
+                content = self._strip_links(content)
+                lines.append(f"[{idx}] Title: {title}\nContent:\n{content}\n")
+            elif type_ == "search":
+                snippet = (res.get("content", "") or "").strip()
+                snippet = self._strip_links(snippet)
+                lines.append(f"[{idx}] Title: {title}\nSnippet: {snippet}\n")
+        
+        return "\n".join(lines)