donkit-llm 0.1.1__py3-none-any.whl → 0.1.2__py3-none-any.whl

donkit/llm/openai_model.py

@@ -21,6 +21,8 @@ from .model_abstract import (
 class OpenAIModel(LLMModelAbstract):
     """OpenAI model implementation supporting GPT-4, GPT-3.5, etc."""
 
+    name = "openai"
+
     def __init__(
         self,
         model_name: str,
@@ -103,13 +105,13 @@ class OpenAIModel(LLMModelAbstract):
             # Multimodal content
             content_parts = []
             for part in msg.content:
-                if part.type == ContentType.TEXT:
+                if part.content_type == ContentType.TEXT:
                     content_parts.append({"type": "text", "text": part.content})
-                elif part.type == ContentType.IMAGE_URL:
+                elif part.content_type == ContentType.IMAGE_URL:
                     content_parts.append(
                         {"type": "image_url", "image_url": {"url": part.content}}
                     )
-                elif part.type == ContentType.IMAGE_BASE64:
+                elif part.content_type == ContentType.IMAGE_BASE64:
                     content_parts.append(
                         {
                             "type": "image_url",
@@ -174,12 +176,19 @@ class OpenAIModel(LLMModelAbstract):
             "messages": messages,
         }
 
-        # if request.temperature is not None:
-        #     kwargs["temperature"] = request.temperature
+        if request.temperature is not None:
+            kwargs["temperature"] = request.temperature
         if request.max_tokens is not None:
-            kwargs["max_completion_tokens"] = (
-                request.max_tokens if request.max_tokens <= 16384 else 16384
-            )
+            # Use max_completion_tokens for GPT models, max_tokens for others
+            model_lower = self._model_name.lower()
+            if "gpt" in model_lower and "oss" not in model_lower:
+                kwargs["max_completion_tokens"] = (
+                    request.max_tokens if request.max_tokens <= 16384 else 16384
+                )
+            else:
+                kwargs["max_tokens"] = (
+                    request.max_tokens if request.max_tokens <= 16384 else 16384
+                )
         if request.top_p is not None:
             kwargs["top_p"] = request.top_p
         if request.stop:
@@ -200,40 +209,47 @@ class OpenAIModel(LLMModelAbstract):
         if request.response_format:
             kwargs["response_format"] = request.response_format
 
-        response = await self.client.chat.completions.create(**kwargs)
+        try:
+            response = await self.client.chat.completions.create(**kwargs)
 
-        choice = response.choices[0]
-        message = choice.message
+            if not response.choices:
+                return GenerateResponse(content="Error: No response choices returned")
 
-        # Extract content
-        content = message.content
+            choice = response.choices[0]
+            message = choice.message
 
-        # Extract tool calls
-        tool_calls = None
-        if message.tool_calls:
-            tool_calls = [
-                ToolCall(
-                    id=tc.id,
-                    type=tc.type,
-                    function=FunctionCall(
-                        name=tc.function.name, arguments=tc.function.arguments
-                    ),
-                )
-                for tc in message.tool_calls
-            ]
+            # Extract content
+            content = message.content
 
-        return GenerateResponse(
-            content=content,
-            tool_calls=tool_calls,
-            finish_reason=choice.finish_reason,
-            usage={
-                "prompt_tokens": response.usage.prompt_tokens,
-                "completion_tokens": response.usage.completion_tokens,
-                "total_tokens": response.usage.total_tokens,
-            }
-            if response.usage
-            else None,
-        )
+            # Extract tool calls
+            tool_calls = None
+            if message.tool_calls:
+                tool_calls = [
+                    ToolCall(
+                        id=tc.id,
+                        type=tc.type,
+                        function=FunctionCall(
+                            name=tc.function.name, arguments=tc.function.arguments
+                        ),
+                    )
+                    for tc in message.tool_calls
+                ]
+
+            return GenerateResponse(
+                content=content,
+                tool_calls=tool_calls,
+                finish_reason=choice.finish_reason,
+                usage={
+                    "prompt_tokens": response.usage.prompt_tokens,
+                    "completion_tokens": response.usage.completion_tokens,
+                    "total_tokens": response.usage.total_tokens,
+                }
+                if response.usage
+                else None,
+            )
+        except Exception as e:
+            error_msg = str(e)
+            return GenerateResponse(content=f"Error: {error_msg}")
 
     async def generate_stream(
         self, request: GenerateRequest
@@ -252,7 +268,16 @@ class OpenAIModel(LLMModelAbstract):
         if request.temperature is not None:
             kwargs["temperature"] = request.temperature
         if request.max_tokens is not None:
-            kwargs["max_tokens"] = request.max_tokens
+            # Use max_completion_tokens for GPT models, max_tokens for others
+            model_lower = self._model_name.lower()
+            if "gpt" in model_lower and "oss" not in model_lower:
+                kwargs["max_completion_tokens"] = (
+                    request.max_tokens if request.max_tokens <= 16384 else 16384
+                )
+            else:
+                kwargs["max_tokens"] = (
+                    request.max_tokens if request.max_tokens <= 16384 else 16384
+                )
         if request.top_p is not None:
             kwargs["top_p"] = request.top_p
         if request.stop:
@@ -273,38 +298,70 @@ class OpenAIModel(LLMModelAbstract):
         if request.response_format:
             kwargs["response_format"] = request.response_format
 
-        stream = await self.client.chat.completions.create(**kwargs)
-
-        async for chunk in stream:
-            if not chunk.choices:
-                continue
-
-            choice = chunk.choices[0]
-            delta = choice.delta
-
-            content = delta.content if delta.content else None
-            finish_reason = choice.finish_reason
-
-            # Handle tool calls in streaming
-            tool_calls = None
-            if delta.tool_calls:
+        try:
+            stream = await self.client.chat.completions.create(**kwargs)
+
+            # Accumulate tool calls across chunks
+            accumulated_tool_calls: dict[int, dict] = {}
+
+            async for chunk in stream:
+                if not chunk.choices:
+                    continue
+
+                choice = chunk.choices[0]
+                delta = choice.delta
+
+                # Yield text content if present
+                if delta.content:
+                    yield StreamChunk(content=delta.content, tool_calls=None)
+
+                # Accumulate tool calls
+                if delta.tool_calls:
+                    for tc_delta in delta.tool_calls:
+                        idx = tc_delta.index
+                        if idx not in accumulated_tool_calls:
+                            accumulated_tool_calls[idx] = {
+                                "id": tc_delta.id or "",
+                                "type": tc_delta.type or "function",
+                                "function": {"name": "", "arguments": ""},
+                            }
+
+                        if tc_delta.id:
+                            accumulated_tool_calls[idx]["id"] = tc_delta.id
+                        if tc_delta.type:
+                            accumulated_tool_calls[idx]["type"] = tc_delta.type
+                        if tc_delta.function:
+                            if tc_delta.function.name:
+                                accumulated_tool_calls[idx]["function"]["name"] = (
+                                    tc_delta.function.name
+                                )
+                            if tc_delta.function.arguments:
+                                accumulated_tool_calls[idx]["function"][
+                                    "arguments"
+                                ] += tc_delta.function.arguments
+
+                # Yield finish reason if present
+                if choice.finish_reason:
+                    yield StreamChunk(content=None, finish_reason=choice.finish_reason)
+
+            # Yield final response with accumulated tool calls if any
+            if accumulated_tool_calls:
                 tool_calls = [
                     ToolCall(
-                        id=tc.id or "",
-                        type=tc.type or "function",
+                        id=tc_data["id"],
+                        type=tc_data["type"],
                         function=FunctionCall(
-                            name=tc.function.name or "",
-                            arguments=tc.function.arguments or "",
+                            name=tc_data["function"]["name"],
+                            arguments=tc_data["function"]["arguments"],
                         ),
                     )
-                    for tc in delta.tool_calls
+                    for tc_data in accumulated_tool_calls.values()
                 ]
+                yield StreamChunk(content=None, tool_calls=tool_calls)
 
-            yield StreamChunk(
-                content=content,
-                tool_calls=tool_calls,
-                finish_reason=finish_reason,
-            )
+        except Exception as e:
+            error_msg = str(e)
+            yield StreamChunk(content=f"Error: {error_msg}")
 
 
 class AzureOpenAIModel(OpenAIModel):
@@ -360,9 +417,8 @@ class AzureOpenAIModel(OpenAIModel):
             | ModelCapability.TOOL_CALLING
             | ModelCapability.STRUCTURED_OUTPUT
             | ModelCapability.MULTIMODAL_INPUT
+            | ModelCapability.VISION
         )
-        if "vision" in self._base_model_name.lower() or "4o" in self._base_model_name:
-            caps |= ModelCapability.MULTIMODAL_INPUT
         return caps
 
     @property
@@ -384,8 +440,27 @@ class AzureOpenAIModel(OpenAIModel):
 
     async def generate(self, request: GenerateRequest) -> GenerateResponse:
         """Generate a response using Azure OpenAI API with parameter adaptation."""
-        # Override to adapt parameters where needed, then call parent
-        return await super().generate(request)
+        # Azure OpenAI uses deployment name instead of model name
+        # Temporarily override model_name with deployment_name
+        original_model = self._model_name
+        self._model_name = self._deployment_name
+        try:
+            return await super().generate(request)
+        finally:
+            self._model_name = original_model
+
+    async def generate_stream(
+        self, request: GenerateRequest
+    ) -> AsyncIterator[StreamChunk]:
+        """Generate a streaming response using Azure OpenAI API."""
+        # Azure OpenAI uses deployment name instead of model name
+        original_model = self._model_name
+        self._model_name = self._deployment_name
+        try:
+            async for chunk in super().generate_stream(request):
+                yield chunk
+        finally:
+            self._model_name = original_model
 
 
 class OpenAIEmbeddingModel(LLMModelAbstract):
@@ -475,19 +550,25 @@ class OpenAIEmbeddingModel(LLMModelAbstract):
         if request.dimensions:
             kwargs["dimensions"] = request.dimensions
 
-        response = await self.client.embeddings.create(**kwargs)
+        try:
+            response = await self.client.embeddings.create(**kwargs)
 
-        embeddings = [item.embedding for item in response.data]
+            embeddings = [item.embedding for item in response.data]
 
-        return EmbeddingResponse(
-            embeddings=embeddings,
-            usage={
-                "prompt_tokens": response.usage.prompt_tokens,
-                "total_tokens": response.usage.total_tokens,
-            }
-            if response.usage
-            else None,
-        )
+            return EmbeddingResponse(
+                embeddings=embeddings,
+                usage={
+                    "prompt_tokens": response.usage.prompt_tokens,
+                    "total_tokens": response.usage.total_tokens,
+                }
+                if response.usage
+                else None,
+                metadata={
+                    "dimensionality": len(embeddings[0]) if len(embeddings) > 0 else 0
+                },
+            )
+        except Exception as e:
+            raise Exception(f"Failed to generate embeddings: {e}")
 
 
 class AzureOpenAIEmbeddingModel(LLMModelAbstract):
@@ -550,12 +631,12 @@ class AzureOpenAIEmbeddingModel(LLMModelAbstract):
         return ModelCapability.EMBEDDINGS
 
     async def generate(self, request: GenerateRequest) -> GenerateResponse:
-        raise NotImplementedError("Embedding models do not support text generation")
+        raise NotImplementedError("Embedding models does not support text generation")
 
     async def generate_stream(
         self, request: GenerateRequest
     ) -> AsyncIterator[StreamChunk]:
-        raise NotImplementedError("Embedding models do not support text generation")
+        raise NotImplementedError("Embedding models does not support text generation")
 
     async def embed(self, request: EmbeddingRequest) -> EmbeddingResponse:
         """Generate embeddings using Azure OpenAI API."""
@@ -566,16 +647,22 @@ class AzureOpenAIEmbeddingModel(LLMModelAbstract):
         if request.dimensions:
             kwargs["dimensions"] = request.dimensions
 
-        response = await self.client.embeddings.create(**kwargs)
+        try:
+            response = await self.client.embeddings.create(**kwargs)
 
-        embeddings = [item.embedding for item in response.data]
+            embeddings = [item.embedding for item in response.data]
 
-        return EmbeddingResponse(
-            embeddings=embeddings,
-            usage={
-                "prompt_tokens": response.usage.prompt_tokens,
-                "total_tokens": response.usage.total_tokens,
-            }
-            if response.usage
-            else None,
-        )
+            return EmbeddingResponse(
+                embeddings=embeddings,
+                usage={
+                    "prompt_tokens": response.usage.prompt_tokens,
+                    "total_tokens": response.usage.total_tokens,
+                }
+                if response.usage
+                else None,
+                metadata={
+                    "dimensions": len(embeddings[0]) if len(embeddings) > 0 else 0
+                },
+            )
+        except Exception as e:
+            raise Exception(f"Failed to generate embeddings: {e}")