isa-model 0.1.1__py3-none-any.whl → 0.2.8__py3-none-any.whl

isa_model/inference/services/llm/openai_llm_service.py

@@ -1,12 +1,13 @@
 import logging
 import os
-from typing import Dict, Any, List, Union, AsyncGenerator, Optional
+import json
+from typing import Dict, Any, List, Union, AsyncGenerator, Optional, Callable
 
 # 使用官方 OpenAI 库和 dotenv
 from openai import AsyncOpenAI
 from dotenv import load_dotenv
 
-from isa_model.inference.services.base_service import BaseLLMService
+from isa_model.inference.services.llm.base_llm_service import BaseLLMService
 from isa_model.inference.providers.base_provider import BaseProvider
 
 # 加载 .env.local 文件中的环境变量
@@ -34,14 +35,85 @@ class OpenAILLMService(BaseLLMService):
             raise ValueError("OPENAI_API_KEY 未设置。") from e
             
         self.last_token_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+        self.total_token_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "requests_count": 0}
+        
+        # Tool binding attributes
+        self._bound_tools: List[Dict[str, Any]] = []
+        self._tool_binding_kwargs: Dict[str, Any] = {}
+        self._tool_functions: Dict[str, Callable] = {}
+        
         logger.info(f"Initialized OpenAILLMService with model {self.model_name} and endpoint {self.client.base_url}")
     
-    async def ainvoke(self, prompt: Union[str, List[Dict[str, str]], Any]) -> str:
+    def _create_bound_copy(self) -> 'OpenAILLMService':
+        """Create a copy of this service for tool binding"""
+        bound_service = OpenAILLMService(self.provider, self.model_name)
+        bound_service._bound_tools = self._bound_tools.copy()
+        bound_service._tool_binding_kwargs = self._tool_binding_kwargs.copy()
+        bound_service._tool_functions = self._tool_functions.copy()
+        return bound_service
+    
+    def bind_tools(self, tools: List[Union[Dict[str, Any], Callable]], **kwargs) -> 'OpenAILLMService':
+        """Bind tools to this LLM service for function calling"""
+        bound_service = self._create_bound_copy()
+        bound_service._bound_tools = self._convert_tools_to_schema(tools)
+        bound_service._tool_binding_kwargs = kwargs
+        
+        # Store the actual functions for execution
+        for tool in tools:
+            if callable(tool):
+                bound_service._tool_functions[tool.__name__] = tool
+        
+        return bound_service
+    
+    async def ainvoke(self, prompt: Union[str, List[Any], Any]) -> str:
         """Universal invocation method"""
         if isinstance(prompt, str):
             return await self.acompletion(prompt)
         elif isinstance(prompt, list):
-            return await self.achat(prompt)
+            # 检查是否是 LangGraph 消息对象
+            if prompt and hasattr(prompt[0], 'content'):
+                # 转换 LangGraph 消息对象为标准格式
+                converted_messages = []
+                for msg in prompt:
+                    if hasattr(msg, 'type'):
+                        # LangGraph 消息对象
+                        msg_dict = {"content": msg.content}
+                        
+                        # 根据消息类型设置 role
+                        if msg.type == "system":
+                            msg_dict["role"] = "system"
+                        elif msg.type == "human":
+                            msg_dict["role"] = "user"
+                        elif msg.type == "ai":
+                            msg_dict["role"] = "assistant"
+                            # 处理工具调用
+                            if hasattr(msg, 'tool_calls') and msg.tool_calls:
+                                msg_dict["tool_calls"] = [
+                                    {
+                                        "id": tc.get("id", f"call_{i}"),
+                                        "type": "function",
+                                        "function": {
+                                            "name": tc["name"],
+                                            "arguments": json.dumps(tc["args"])
+                                        }
+                                    } for i, tc in enumerate(msg.tool_calls)
+                                ]
+                        elif msg.type == "tool":
+                            msg_dict["role"] = "tool"
+                            if hasattr(msg, 'tool_call_id'):
+                                msg_dict["tool_call_id"] = msg.tool_call_id
+                        else:
+                            msg_dict["role"] = "user"  # 默认为用户消息
+                        
+                        converted_messages.append(msg_dict)
+                    else:
+                        # 已经是字典格式
+                        converted_messages.append(msg)
+                
+                return await self.achat(converted_messages)
+            else:
+                # 标准字典格式的消息
+                return await self.achat(prompt)
         else:
             raise ValueError("Prompt must be a string or a list of messages")
     
@@ -51,12 +123,19 @@ class OpenAILLMService(BaseLLMService):
             temperature = self.config.get("temperature", 0.7)
             max_tokens = self.config.get("max_tokens", 1024)
             
-            response = await self.client.chat.completions.create(
-                model=self.model_name,
-                messages=messages,
-                temperature=temperature,
-                max_tokens=max_tokens
-            )
+            kwargs = {
+                "model": self.model_name,
+                "messages": messages,
+                "temperature": temperature,
+                "max_tokens": max_tokens
+            }
+            
+            # Add tools if bound
+            if self._has_bound_tools():
+                kwargs["tools"] = self._get_bound_tools()
+                kwargs["tool_choice"] = "auto"
+            
+            response = await self.client.chat.completions.create(**kwargs)
             
             if response.usage:
                 self.last_token_usage = {
@@ -64,13 +143,87 @@ class OpenAILLMService(BaseLLMService):
                     "completion_tokens": response.usage.completion_tokens,
                     "total_tokens": response.usage.total_tokens
                 }
+                
+                # Update total usage
+                self.total_token_usage["prompt_tokens"] += self.last_token_usage["prompt_tokens"]
+                self.total_token_usage["completion_tokens"] += self.last_token_usage["completion_tokens"]
+                self.total_token_usage["total_tokens"] += self.last_token_usage["total_tokens"]
+                self.total_token_usage["requests_count"] += 1
             
-            return response.choices[0].message.content or ""
+            # Handle tool calls if present
+            message = response.choices[0].message
+            if message.tool_calls:
+                return await self._handle_tool_calls(message, messages)
+            
+            return message.content or ""
             
         except Exception as e:
             logger.error(f"Error in chat completion: {e}")
             raise
     
+    async def _handle_tool_calls(self, assistant_message, original_messages: List[Dict[str, str]]) -> str:
+        """Handle tool calls from the assistant"""
+        # Add assistant message with tool calls to conversation
+        messages = original_messages + [{
+            "role": "assistant", 
+            "content": assistant_message.content or "",
+            "tool_calls": [
+                {
+                    "id": tc.id,
+                    "type": tc.type,
+                    "function": {
+                        "name": tc.function.name,
+                        "arguments": tc.function.arguments
+                    }
+                } for tc in assistant_message.tool_calls
+            ]
+        }]
+        
+        # Execute each tool call
+        for tool_call in assistant_message.tool_calls:
+            function_name = tool_call.function.name
+            arguments = json.loads(tool_call.function.arguments)
+            
+            try:
+                # Execute the tool
+                if function_name in self._tool_functions:
+                    result = self._tool_functions[function_name](**arguments)
+                    if hasattr(result, '__await__'):  # Handle async functions
+                        result = await result
+                else:
+                    result = f"Error: Function {function_name} not found"
+                
+                # Add tool result to messages
+                messages.append({
+                    "role": "tool",
+                    "content": str(result),
+                    "tool_call_id": tool_call.id
+                })
+                
+            except Exception as e:
+                logger.error(f"Error executing tool {function_name}: {e}")
+                messages.append({
+                    "role": "tool",
+                    "content": f"Error executing {function_name}: {str(e)}",
+                    "tool_call_id": tool_call.id
+                })
+        
+        # Get final response from the model with all context
+        try:
+            kwargs = {
+                "model": self.model_name,
+                "messages": messages,
+                "temperature": self.config.get("temperature", 0.7),
+                "max_tokens": self.config.get("max_tokens", 1024)
+            }
+            
+            response = await self.client.chat.completions.create(**kwargs)
+            return response.choices[0].message.content or ""
+            
+        except Exception as e:
+            logger.error(f"Error getting final response after tool calls: {e}")
+            raise
+    
     async def acompletion(self, prompt: str) -> str:
         """Text completion method (using chat API)"""
         messages = [{"role": "user", "content": prompt}]
@@ -82,13 +235,20 @@ class OpenAILLMService(BaseLLMService):
             temperature = self.config.get("temperature", 0.7)
             max_tokens = self.config.get("max_tokens", 1024)
             
-            response = await self.client.chat.completions.create(
-                model=self.model_name,
-                messages=messages,
-                temperature=temperature,
-                max_tokens=max_tokens,
-                n=n
-            )
+            kwargs = {
+                "model": self.model_name,
+                "messages": messages,
+                "temperature": temperature,
+                "max_tokens": max_tokens,
+                "n": n
+            }
+            
+            # Add tools if bound
+            if self._has_bound_tools():
+                kwargs["tools"] = self._get_bound_tools()
+                kwargs["tool_choice"] = "auto"
+            
+            response = await self.client.chat.completions.create(**kwargs)
             
             if response.usage:
                 self.last_token_usage = {
@@ -96,6 +256,12 @@ class OpenAILLMService(BaseLLMService):
                     "completion_tokens": response.usage.completion_tokens,
                     "total_tokens": response.usage.total_tokens
                 }
+                
+                # Update total usage
+                self.total_token_usage["prompt_tokens"] += self.last_token_usage["prompt_tokens"]
+                self.total_token_usage["completion_tokens"] += self.last_token_usage["completion_tokens"]
+                self.total_token_usage["total_tokens"] += self.last_token_usage["total_tokens"]
+                self.total_token_usage["requests_count"] += 1
             
             return [choice.message.content or "" for choice in response.choices]
         except Exception as e:
@@ -108,13 +274,20 @@ class OpenAILLMService(BaseLLMService):
             temperature = self.config.get("temperature", 0.7)
             max_tokens = self.config.get("max_tokens", 1024)
             
-            stream = await self.client.chat.completions.create(
-                model=self.model_name,
-                messages=messages,
-                temperature=temperature,
-                max_tokens=max_tokens,
-                stream=True
-            )
+            kwargs = {
+                "model": self.model_name,
+                "messages": messages,
+                "temperature": temperature,
+                "max_tokens": max_tokens,
+                "stream": True
+            }
+            
+            # Add tools if bound
+            if self._has_bound_tools():
+                kwargs["tools"] = self._get_bound_tools()
+                kwargs["tool_choice"] = "auto"
+            
+            stream = await self.client.chat.completions.create(**kwargs)
             
             async for chunk in stream:
                 content = chunk.choices[0].delta.content
@@ -125,14 +298,38 @@ class OpenAILLMService(BaseLLMService):
             logger.error(f"Error in stream chat: {e}")
             raise
     
-    def get_token_usage(self) -> Dict[str, int]:
+    async def astream_completion(self, prompt: str) -> AsyncGenerator[str, None]:
+        """Stream completion responses"""
+        messages = [{"role": "user", "content": prompt}]
+        async for chunk in self.astream_chat(messages):
+            yield chunk
+    
+    def get_token_usage(self) -> Dict[str, Any]:
         """Get total token usage statistics"""
-        return self.last_token_usage
+        return self.total_token_usage
     
     def get_last_token_usage(self) -> Dict[str, int]:
         """Get token usage from last request"""
         return self.last_token_usage
+    
+    def get_model_info(self) -> Dict[str, Any]:
+        """Get information about the current model"""
+        return {
+            "name": self.model_name,
+            "max_tokens": self.config.get("max_tokens", 1024),
+            "supports_streaming": True,
+            "supports_functions": True,
+            "provider": "openai"
+        }
+    
+    def _has_bound_tools(self) -> bool:
+        """Check if this service has bound tools"""
+        return bool(self._bound_tools)
+    
+    def _get_bound_tools(self) -> List[Dict[str, Any]]:
+        """Get the bound tools schema"""
+        return self._bound_tools
         
     async def close(self):
         """Close the backend client"""
-        await self.client.aclose()
+        await self.client.close()