isa-model 0.4.0__py3-none-any.whl → 0.4.4__py3-none-any.whl

isa_model/inference/services/llm/base_llm_service.py

@@ -1,9 +1,12 @@
 from abc import ABC, abstractmethod
 from typing import Dict, Any, List, Union, Optional, AsyncGenerator, Callable
 import logging
+import json
 
 from isa_model.inference.services.base_service import BaseService
 from isa_model.inference.services.llm.helpers.llm_adapter import AdapterManager
+from isa_model.inference.services.llm.helpers.llm_utils import TokenCounter, TextProcessor, ResponseParser, LLMMetrics
+from isa_model.inference.services.llm.helpers.llm_prompts import LLMPrompts, LLMPromptTemplates
 
 logger = logging.getLogger(__name__)
 
@@ -18,6 +21,12 @@ class BaseLLMService(BaseService):
         # 初始化适配器管理器
         self.adapter_manager = AdapterManager()
         
+        # Initialize helper utilities (optional, can be overridden by specific services)
+        self.token_counter = TokenCounter(model_name)
+        self.text_processor = TextProcessor()
+        self.response_parser = ResponseParser()
+        self.llm_prompts = LLMPrompts()
+        
         # Get config from provider
         provider_config = self.get_provider_config()
         self.streaming = provider_config.get("streaming", False)
@@ -29,6 +38,8 @@ class BaseLLMService(BaseService):
         input_data: Union[str, List[Dict[str, str]], Any],
         task: Optional[str] = None,
         show_reasoning: bool = False,
+        output_format: Optional[str] = None,
+        json_schema: Optional[Dict] = None,
         **kwargs
     ) -> Dict[str, Any]:
         """
@@ -40,75 +51,186 @@ class BaseLLMService(BaseService):
                 - list: 消息历史 [{"role": "user", "content": "hello"}]
                 - Any: LangChain 消息对象或其他格式
             task: 任务类型，支持多种LLM任务
+            output_format: Output format ("json", "markdown", "code", etc.)
+            json_schema: JSON schema for structured output validation
             **kwargs: 任务特定的附加参数
             
         Returns:
-            Dict containing task results
+            Dict containing task results (optionally formatted as JSON)
         """
         task = task or "chat"
         
+        # Store formatting options for use by specific task methods
+        format_options = {
+            "output_format": output_format,
+            "json_schema": json_schema,
+            "repair_attempts": kwargs.get("repair_attempts", 3)
+        }
+        
+        # Execute task and apply formatting
+        result = None
+        
         # ==================== 对话类任务 ====================
         if task == "chat":
-            return await self.chat(input_data, kwargs.get("max_tokens", self.max_tokens), show_reasoning=show_reasoning)
+            # Pass all kwargs to ainvoke for better parameter support (like response_format)
+            result_raw = await self.ainvoke(input_data, show_reasoning=show_reasoning, **kwargs)
+            # Wrap in chat response format, preserving AIMessage objects with tool_calls
+            if hasattr(result_raw, 'tool_calls'):
+                # This is an AIMessage with tool_calls - preserve the entire object
+                result = {"message": result_raw}
+            elif hasattr(result_raw, 'content'):
+                # Regular AIMessage without tool_calls - extract content
+                content = result_raw.content
+                result = {"message": content}
+            else:
+                # Plain string response
+                content = str(result_raw)
+                result = {"message": content}
         elif task == "complete":
-            return await self.complete_text(input_data, kwargs.get("max_tokens", self.max_tokens))
+            result = await self.complete_text(input_data, kwargs.get("max_tokens", self.max_tokens))
         elif task == "instruct":
-            return await self.instruct(input_data, kwargs.get("instruction"), kwargs.get("max_tokens", self.max_tokens))
+            result = await self.instruct(input_data, kwargs.get("instruction"), kwargs.get("max_tokens", self.max_tokens))
         
         # ==================== 文本生成类任务 ====================
         elif task == "generate":
-            return await self.generate_text(input_data, kwargs.get("max_tokens", self.max_tokens))
+            result = await self.generate_text(input_data, kwargs.get("max_tokens", self.max_tokens))
         elif task == "rewrite":
-            return await self.rewrite_text(input_data, kwargs.get("style"), kwargs.get("tone"))
+            result = await self.rewrite_text(input_data, kwargs.get("style"), kwargs.get("tone"))
         elif task == "summarize":
-            return await self.summarize_text(input_data, kwargs.get("max_length"), kwargs.get("style"))
+            result = await self.summarize_text(input_data, kwargs.get("max_length"), kwargs.get("style"))
         elif task == "translate":
             target_language = kwargs.get("target_language")
             if not target_language:
                 raise ValueError("target_language is required for translate task")
-            return await self.translate_text(input_data, target_language, kwargs.get("source_language"))
+            result = await self.translate_text(input_data, target_language, kwargs.get("source_language"))
         
         # ==================== 分析类任务 ====================
         elif task == "analyze":
-            return await self.analyze_text(input_data, kwargs.get("analysis_type"))
+            result = await self.analyze_text(input_data, kwargs.get("analysis_type"))
         elif task == "classify":
-            return await self.classify_text(input_data, kwargs.get("categories"))
+            result = await self.classify_text(input_data, kwargs.get("categories"))
         elif task == "extract":
-            return await self.extract_information(input_data, kwargs.get("extract_type"))
+            result = await self.extract_information(input_data, kwargs.get("extract_type"))
         elif task == "sentiment":
-            return await self.analyze_sentiment(input_data)
+            # Always use chat with appropriate prompt for sentiment analysis
+            if output_format == "json":
+                # Create JSON-formatted prompt
+                json_prompt = self.create_json_prompt(
+                    f"Please analyze the sentiment of the following text: {input_data}",
+                    json_schema or {
+                        "type": "object",
+                        "properties": {
+                            "sentiment": {"type": "string", "enum": ["positive", "negative", "neutral"]},
+                            "confidence": {"type": "number", "minimum": 0, "maximum": 1},
+                            "explanation": {"type": "string"}
+                        },
+                        "required": ["sentiment"]
+                    }
+                )
+                result = await self.chat(json_prompt, show_reasoning=show_reasoning)
+            else:
+                # Use simple chat prompt for sentiment analysis
+                sentiment_prompt = f"Please analyze the sentiment of the following text and classify it as positive, negative, or neutral:\n\n{input_data}\n\nSentiment:"
+                result = await self.chat(sentiment_prompt, show_reasoning=show_reasoning)
         
         # ==================== 编程类任务 ====================
         elif task == "code":
-            return await self.generate_code(input_data, kwargs.get("language"), kwargs.get("style"))
+            # Always use chat with appropriate prompt for code generation
+            language = kwargs.get("language", "")
+            style = kwargs.get("style", "")
+            
+            code_prompt = f"Please write code"
+            if language:
+                code_prompt += f" in {language}"
+            code_prompt += f" for the following requirement:\n\n{input_data}\n\n"
+            
+            if style:
+                code_prompt += f"Style requirements: {style}\n\n"
+            
+            code_prompt += "Please provide clean, working code with comments."
+            
+            result = await self.chat(code_prompt, show_reasoning=show_reasoning)
         elif task == "explain_code":
-            return await self.explain_code(input_data, kwargs.get("language"))
+            result = await self.explain_code(input_data, kwargs.get("language"))
         elif task == "debug_code":
-            return await self.debug_code(input_data, kwargs.get("language"))
+            result = await self.debug_code(input_data, kwargs.get("language"))
         elif task == "refactor_code":
-            return await self.refactor_code(input_data, kwargs.get("language"), kwargs.get("improvements"))
+            result = await self.refactor_code(input_data, kwargs.get("language"), kwargs.get("improvements"))
         
         # ==================== 推理类任务 ====================
         elif task == "reason":
-            return await self.reason_about(input_data, kwargs.get("reasoning_type"))
+            # Always use chat with appropriate prompt for reasoning
+            reasoning_type = kwargs.get("reasoning_type", "")
+            
+            reason_prompt = f"Please analyze and explain the reasoning behind the following question or topic"
+            if reasoning_type:
+                reason_prompt += f" using {reasoning_type} reasoning"
+            reason_prompt += f":\n\n{input_data}\n\n"
+            reason_prompt += "Provide a clear, step-by-step explanation of your reasoning process."
+            
+            result = await self.chat(reason_prompt, show_reasoning=show_reasoning)
         elif task == "solve":
-            return await self.solve_problem(input_data, kwargs.get("problem_type"))
+            # Always use chat with appropriate prompt for problem solving
+            problem_type = kwargs.get("problem_type", "")
+            
+            solve_prompt = f"Please solve the following problem"
+            if problem_type:
+                solve_prompt += f" (type: {problem_type})"
+            solve_prompt += f":\n\n{input_data}\n\n"
+            solve_prompt += "Provide a clear solution with step-by-step explanation."
+            
+            result = await self.chat(solve_prompt, show_reasoning=show_reasoning)
         elif task == "plan":
-            return await self.create_plan(input_data, kwargs.get("plan_type"))
+            result = await self.create_plan(input_data, kwargs.get("plan_type"))
         elif task == "deep_research":
-            return await self.deep_research(input_data, kwargs.get("research_type"), kwargs.get("search_enabled", True))
+            result = await self.deep_research(input_data, kwargs.get("research_type"), kwargs.get("search_enabled", True))
         
         # ==================== 工具调用类任务 ====================
         elif task == "tool_call":
-            return await self.call_tools(input_data, kwargs.get("available_tools"))
+            result = await self.call_tools(input_data, kwargs.get("available_tools"))
         elif task == "function_call":
             function_name = kwargs.get("function_name")
             if not function_name:
                 raise ValueError("function_name is required for function_call task")
-            return await self.call_function(input_data, function_name, kwargs.get("parameters"))
+            result = await self.call_function(input_data, function_name, kwargs.get("parameters"))
         
         else:
             raise NotImplementedError(f"{self.__class__.__name__} does not support task: {task}")
+        
+        # Apply output formatting if requested
+        if result is not None and output_format:
+            # Extract the raw response for formatting
+            # If result is a dict with 'message' key, use the message for formatting
+            format_input = result
+            if isinstance(result, dict) and 'message' in result:
+                format_input = result['message']
+            
+            formatted_result = self.format_structured_output(
+                response=format_input,
+                output_format=output_format,
+                schema=json_schema,
+                repair_attempts=format_options.get("repair_attempts", 3)
+            )
+            
+            # If formatting succeeded, return formatted result
+            if formatted_result.get("success", False):
+                return {
+                    "result": formatted_result["data"],
+                    "formatted": True,
+                    "format": output_format,
+                    "original": result
+                }
+            else:
+                # If formatting failed, return original with error info
+                return {
+                    "result": result,
+                    "formatted": False,
+                    "format_errors": formatted_result.get("errors", []),
+                    "original": result
+                }
+        
+        # Return unformatted result
+        return result if result is not None else {"message": "Task completed but returned no result"}
     
     # ==================== 对话类方法 ====================
     
@@ -119,7 +241,7 @@ class BaseLLMService(BaseService):
         show_reasoning: bool = False
     ) -> Dict[str, Any]:
         """
-        对话聊天 - Provider必须实现
+        对话聊天 - 委托给 ainvoke 方法
         
         Args:
             input_data: 输入消息
@@ -129,7 +251,19 @@ class BaseLLMService(BaseService):
         Returns:
             Dict containing chat response
         """
-        raise NotImplementedError(f"{self.__class__.__name__} does not support chat task")
+        result = await self.ainvoke(input_data, show_reasoning=show_reasoning)
+        # Ensure we return a proper response structure
+        if result is None:
+            logger.warning("ainvoke returned None - this may indicate an implementation issue")
+            return {"message": ""}
+        
+        # Extract content if it's an AIMessage object
+        if hasattr(result, 'content'):
+            content = result.content
+        else:
+            content = str(result)
+            
+        return {"message": content}
     
     # ==================== 文本生成类方法 ====================
     
@@ -514,6 +648,183 @@ class BaseLLMService(BaseService):
             )
         return 0.0
     
+    # ==================== JSON OUTPUT AND FORMATTING METHODS ====================
+    
+    def format_structured_output(
+        self, 
+        response: Union[str, Any], 
+        output_format: str = "json",
+        schema: Optional[Dict] = None,
+        repair_attempts: int = 3
+    ) -> Dict[str, Any]:
+        """
+        Format response as structured output (JSON, etc.)
+        
+        Args:
+            response: Raw response from model
+            output_format: Desired output format ("json", "code", "structured")
+            schema: Optional JSON schema for validation
+            repair_attempts: Number of JSON repair attempts
+            
+        Returns:
+            Dict with formatted output and metadata
+        """
+        if output_format == "json":
+            if isinstance(response, str):
+                return self.text_processor.extract_json_from_text(response, schema, repair_attempts)
+            else:
+                # Handle response objects with content attribute
+                content = getattr(response, 'content', str(response))
+                return self.text_processor.extract_json_from_text(content, schema, repair_attempts)
+        
+        elif output_format == "code":
+            content = response if isinstance(response, str) else getattr(response, 'content', str(response))
+            code_blocks = self.text_processor.extract_code_blocks(content)
+            return {
+                "success": True,
+                "data": code_blocks,
+                "method": "code_block_extraction",
+                "errors": []
+            }
+        
+        elif output_format == "structured":
+            # Use ResponseParser for general structured parsing
+            content = response if isinstance(response, str) else getattr(response, 'content', str(response))
+            parsed = self.response_parser.parse_structured_response(content, "json")
+            if parsed:
+                return {
+                    "success": True,
+                    "data": parsed,
+                    "method": "structured_parsing",
+                    "errors": []
+                }
+            else:
+                return {
+                    "success": False,
+                    "data": content,
+                    "method": "raw_fallback",
+                    "errors": ["Failed to parse as structured output"]
+                }
+        
+        # Fallback: return raw response
+        return {
+            "success": True,
+            "data": response,
+            "method": "raw_output",
+            "errors": []
+        }
+    
+    def create_json_prompt(
+        self, 
+        base_prompt: str, 
+        json_schema: Optional[Dict] = None,
+        output_instructions: Optional[str] = None
+    ) -> str:
+        """
+        Create a prompt that requests JSON output
+        
+        Args:
+            base_prompt: The base prompt content
+            json_schema: Optional JSON schema to include in prompt
+            output_instructions: Custom output format instructions
+            
+        Returns:
+            Enhanced prompt requesting JSON output
+        """
+        if output_instructions:
+            json_instruction = output_instructions
+        else:
+            json_instruction = LLMPromptTemplates.OUTPUT_FORMATS["json"]
+        
+        if json_schema:
+            schema_text = f"\n\nPlease format your response according to this JSON schema:\n```json\n{json.dumps(json_schema, indent=2)}\n```"
+            return f"{base_prompt}{schema_text}\n\n{json_instruction}"
+        else:
+            return f"{base_prompt}\n\n{json_instruction}"
+    
+    def create_structured_prompt(
+        self, 
+        task_type: str,
+        content: str,
+        output_format: str = "json",
+        **kwargs
+    ) -> str:
+        """
+        Create a structured prompt using LLMPrompts templates
+        
+        Args:
+            task_type: Type of task (from LLMPrompts methods)
+            content: Main content/input
+            output_format: Desired output format
+            **kwargs: Additional arguments for the prompt template
+            
+        Returns:
+            Formatted prompt string
+        """
+        try:
+            # Get the appropriate prompt template
+            if hasattr(self.llm_prompts, f"{task_type}_prompt"):
+                method = getattr(self.llm_prompts, f"{task_type}_prompt")
+                base_prompt = method(content, **kwargs)
+            else:
+                # Fallback to generic prompt
+                base_prompt = f"Please {task_type} the following:\n\n{content}"
+            
+            # Add output format instructions
+            if output_format in LLMPromptTemplates.OUTPUT_FORMATS:
+                format_instruction = LLMPromptTemplates.OUTPUT_FORMATS[output_format]
+                return f"{base_prompt}\n\n{format_instruction}"
+            
+            return base_prompt
+            
+        except Exception as e:
+            logger.warning(f"Failed to create structured prompt: {e}")
+            return f"Please {task_type} the following:\n\n{content}"
+    
+    def count_tokens(self, text: Union[str, List[Dict[str, str]]]) -> int:
+        """
+        Count tokens in text or message list
+        
+        Args:
+            text: String or message list to count tokens for
+            
+        Returns:
+            Number of tokens
+        """
+        if isinstance(text, str):
+            return self.token_counter.count_tokens(text)
+        elif isinstance(text, list):
+            return self.token_counter.count_messages_tokens(text)
+        else:
+            return self.token_counter.count_tokens(str(text))
+    
+    def truncate_to_token_limit(self, text: str, max_tokens: int) -> str:
+        """
+        Truncate text to fit within token limit
+        
+        Args:
+            text: Text to truncate
+            max_tokens: Maximum number of tokens
+            
+        Returns:
+            Truncated text
+        """
+        return self.token_counter.truncate_text(text, max_tokens)
+    
+    def split_text_by_tokens(self, text: str, chunk_size: int, overlap: int = 0) -> List[str]:
+        """
+        Split text into chunks by token count
+        
+        Args:
+            text: Text to split
+            chunk_size: Size of each chunk in tokens
+            overlap: Number of overlapping tokens between chunks
+            
+        Returns:
+            List of text chunks
+        """
+        return self.token_counter.split_text_by_tokens(text, chunk_size, overlap)
+
     # ==================== METADATA AND UTILITY METHODS ====================
     
     def get_supported_tasks(self) -> List[str]: