isa-model 0.0.3__py3-none-any.whl → 0.0.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,8 +35,36 @@ 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}")
     
+    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[Dict[str, str]], Any]) -> str:
         """Universal invocation method"""
         if isinstance(prompt, str):
@@ -51,12 +80,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 +100,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 +192,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 +213,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 +231,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 +255,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()