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

donkit/llm/__init__.py

@@ -23,6 +23,8 @@ from .openai_model import (
 from .claude_model import ClaudeModel, ClaudeVertexModel
 from .vertex_model import VertexAIModel, VertexEmbeddingModel
 from .factory import ModelFactory
+from .gemini_model import GeminiModel, GeminiEmbeddingModel
+from .donkit_model import DonkitModel
 
 __all__ = [
     "ModelFactory",
@@ -52,4 +54,7 @@ __all__ = [
     "ClaudeVertexModel",
     "VertexAIModel",
     "VertexEmbeddingModel",
+    "GeminiModel",
+    "GeminiEmbeddingModel",
+    "DonkitModel",
 ]

donkit/llm/claude_model.py

@@ -20,6 +20,8 @@ from .model_abstract import (
 class ClaudeModel(LLMModelAbstract):
     """Anthropic Claude model implementation."""
 
+    name = "claude"
+
     def __init__(
         self,
         model_name: str,
@@ -82,9 +84,9 @@ class ClaudeModel(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:
                     # Claude expects base64 images, not URLs
                     content_parts.append(
                         {
@@ -95,7 +97,7 @@ class ClaudeModel(LLMModelAbstract):
                             },
                         }
                     )
-                elif part.type == ContentType.IMAGE_BASE64:
+                elif part.content_type == ContentType.IMAGE_BASE64:
                     content_parts.append(
                         {
                             "type": "image",
@@ -319,9 +321,9 @@ class ClaudeVertexModel(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_BASE64:
+                elif part.content_type == ContentType.IMAGE_BASE64:
                     content_parts.append(
                         {
                             "type": "image",

donkit/llm/donkit_model.py

@@ -0,0 +1,238 @@
+from typing import Any, AsyncIterator
+
+from donkit.ragops_api_gateway_client.client import RagopsAPIGatewayClient
+from .model_abstract import (
+    EmbeddingRequest,
+    EmbeddingResponse,
+    FunctionCall,
+    GenerateRequest,
+    GenerateResponse,
+    LLMModelAbstract,
+    Message,
+    ModelCapability,
+    StreamChunk,
+    Tool,
+    ToolCall,
+)
+
+
+class DonkitModel(LLMModelAbstract):
+    """
+    Implementation of LLMModelAbstract that proxies requests via RagopsAPIGatewayClient.
+    """
+
+    name = "donkit"
+
+    def __init__(
+        self,
+        base_url: str,
+        api_token: str,
+        provider: str,
+        model_name: str | None = None,
+        project_id: str | None = None,
+    ):
+        """
+        Initialize DonkitModel.
+
+        Args:
+            base_url: Base URL for the API Gateway
+            api_token: API token for authentication
+            provider: The LLM provider name (e.g., "openai", "anthropic", "vertex")
+            model_name: The specific model identifier (e.g., "gpt-4o", "claude-3-opus")
+            project_id: The project ID for the gateway
+        """
+        self.base_url = base_url
+        self.api_token = api_token
+        self.provider = provider
+        self._model_name = model_name
+        self.project_id = project_id
+        self._capabilities = self._determine_capabilities()
+
+    @property
+    def model_name(self) -> str:
+        return self._model_name
+
+    @model_name.setter
+    def model_name(self, value: str):
+        self._model_name = value
+        self._capabilities = self._determine_capabilities()
+
+    @property
+    def capabilities(self) -> ModelCapability:
+        return self._capabilities
+
+    def _determine_capabilities(self) -> ModelCapability:
+        """
+        Estimate capabilities based on model name.
+        Since this is a proxy, we assume modern defaults but refine based on keywords.
+        """
+        caps = (
+            ModelCapability.TEXT_GENERATION
+            | ModelCapability.STREAMING
+            | ModelCapability.STRUCTURED_OUTPUT
+            | ModelCapability.TOOL_CALLING
+            | ModelCapability.MULTIMODAL_INPUT
+            | ModelCapability.EMBEDDINGS
+        )
+        return caps
+
+    def _convert_message(self, msg: Message) -> dict:
+        """Convert internal Message to dictionary format expected by the Gateway."""
+        result: dict[str, Any] = {"role": msg.role}
+        if isinstance(msg.content, str):
+            result["content"] = msg.content
+        else:
+            # Multimodal content processing
+            content_parts = []
+            for part in msg.content if msg.content else []:
+                content_parts.append(part.model_dump(exclude_none=True))
+            result["content"] = content_parts
+        if msg.tool_calls:
+            result["tool_calls"] = [tc.model_dump() for tc in msg.tool_calls]
+        if msg.tool_call_id:
+            result["tool_call_id"] = msg.tool_call_id
+        if msg.name:
+            result["name"] = msg.name
+
+        return result
+
+    def _convert_tools(self, tools: list[Tool]) -> list[dict]:
+        """Convert internal Tool definitions to Gateway dictionary format."""
+        return [tool.model_dump(exclude_none=True) for tool in tools]
+
+    def _prepare_generate_kwargs(self, request: GenerateRequest) -> dict:
+        """Prepare kwargs for generate/generate_stream calls."""
+        messages = [self._convert_message(msg) for msg in request.messages]
+        tools_payload = self._convert_tools(request.tools) if request.tools else None
+
+        kwargs: dict[str, Any] = {
+            "provider": self.provider,
+            "model_name": self._model_name,
+            "messages": messages,
+            "project_id": self.project_id,
+        }
+
+        if request.temperature is not None:
+            kwargs["temperature"] = request.temperature
+        if request.max_tokens is not None:
+            kwargs["max_tokens"] = request.max_tokens
+        if request.top_p is not None:
+            kwargs["top_p"] = request.top_p
+        if request.stop:
+            kwargs["stop"] = request.stop
+        if tools_payload:
+            kwargs["tools"] = tools_payload
+            if request.tool_choice:
+                if isinstance(request.tool_choice, (str, dict)):
+                    kwargs["tool_choice"] = request.tool_choice
+                else:
+                    kwargs["tool_choice"] = "auto"
+        if request.response_format:
+            kwargs["response_format"] = request.response_format
+
+        return kwargs
+
+    async def generate(self, request: GenerateRequest) -> GenerateResponse:
+        """Generate a response using RagopsAPIGatewayClient."""
+        await self.validate_request(request)
+
+        kwargs = self._prepare_generate_kwargs(request)
+
+        async with RagopsAPIGatewayClient(
+            base_url=self.base_url,
+            api_token=self.api_token,
+        ) as client:
+            response_dict = await client.generate(**kwargs)
+
+        # Gateway returns simplified format: {content, tool_calls, finish_reason, usage}
+        content = response_dict.get("content")
+        finish_reason = response_dict.get("finish_reason")
+
+        # Extract tool calls
+        tool_calls = None
+        if response_dict.get("tool_calls"):
+            tool_calls = [
+                ToolCall(
+                    id=tc.get("id"),
+                    type=tc.get("type", "function"),
+                    function=FunctionCall(
+                        name=tc.get("function", {}).get("name"),
+                        arguments=tc.get("function", {}).get("arguments"),
+                    ),
+                )
+                for tc in response_dict["tool_calls"]
+            ]
+
+        usage_data = response_dict.get("usage", {})
+
+        return GenerateResponse(
+            content=content,
+            tool_calls=tool_calls,
+            finish_reason=finish_reason,
+            usage={
+                "prompt_tokens": usage_data.get("prompt_tokens"),
+                "completion_tokens": usage_data.get("completion_tokens"),
+                "total_tokens": usage_data.get("total_tokens"),
+            }
+            if usage_data
+            else None,
+            metadata=response_dict.get("metadata"),
+        )
+
+    async def generate_stream(
+        self, request: GenerateRequest
+    ) -> AsyncIterator[StreamChunk]:
+        """Generate a streaming response using RagopsAPIGatewayClient."""
+        await self.validate_request(request)
+
+        kwargs = self._prepare_generate_kwargs(request)
+
+        async with RagopsAPIGatewayClient(
+            base_url=self.base_url,
+            api_token=self.api_token,
+        ) as client:
+            # Iterate over the stream from client
+            async for chunk_dict in client.generate_stream(**kwargs):
+                content = chunk_dict.get("content")
+                finish_reason = chunk_dict.get("finish_reason")
+
+                tool_calls = None
+                if chunk_dict.get("tool_calls"):
+                    tool_calls = [
+                        ToolCall(
+                            id=tc.get("id", ""),
+                            type=tc.get("type", "function"),
+                            function=FunctionCall(
+                                name=tc.get("function", {}).get("name", ""),
+                                arguments=tc.get("function", {}).get("arguments", ""),
+                            ),
+                        )
+                        for tc in chunk_dict["tool_calls"]
+                    ]
+
+                yield StreamChunk(
+                    content=content,
+                    tool_calls=tool_calls,
+                    finish_reason=finish_reason,
+                    metadata=chunk_dict.get("metadata", {}),
+                )
+
+    async def embed(self, request: EmbeddingRequest) -> EmbeddingResponse:
+        """Generate embeddings using RagopsAPIGatewayClient."""
+
+        kwargs: dict[str, Any] = {
+            "provider": self.provider,
+            "input": request.input,
+            "model_name": self._model_name,
+            "project_id": self.project_id,
+        }
+
+        if request.dimensions:
+            kwargs["dimensions"] = request.dimensions
+        async with RagopsAPIGatewayClient(
+            base_url=self.base_url,
+            api_token=self.api_token,
+        ) as client:
+            response_dict = await client.embeddings(**kwargs)
+
+        return EmbeddingResponse(**response_dict)

donkit/llm/factory.py

@@ -1,14 +1,16 @@
 from typing import Literal
 
-from .claude_model import ClaudeModel, ClaudeVertexModel
+from .claude_model import ClaudeModel
+from .claude_model import ClaudeVertexModel
+from .donkit_model import DonkitModel
+from .gemini_model import GeminiModel
 from .model_abstract import LLMModelAbstract
-from .openai_model import (
-    AzureOpenAIEmbeddingModel,
-    AzureOpenAIModel,
-    OpenAIEmbeddingModel,
-    OpenAIModel,
-)
-from .vertex_model import VertexAIModel, VertexEmbeddingModel
+from .openai_model import AzureOpenAIEmbeddingModel
+from .openai_model import AzureOpenAIModel
+from .openai_model import OpenAIEmbeddingModel
+from .openai_model import OpenAIModel
+from .vertex_model import VertexAIModel
+from .vertex_model import VertexEmbeddingModel
 
 
 class ModelFactory:
@@ -46,7 +48,7 @@ class ModelFactory:
 
     @staticmethod
     def create_embedding_model(
-        provider: Literal["openai", "azure_openai", "vertex"],
+        provider: Literal["openai", "azure_openai", "vertex", "custom", "default"],
         model_name: str | None = None,
         api_key: str | None = None,
         **kwargs,
@@ -92,6 +94,35 @@ class ModelFactory:
             base_url=base_url,
         )
 
+    @staticmethod
+    def create_gemini_model(
+        model_name: str,
+        api_key: str | None = None,
+        project_id: str | None = None,
+        location: str = "us-central1",
+        use_vertex: bool = False,
+    ) -> GeminiModel:
+        """
+        Create a Gemini model instance.
+
+        Args:
+            model_name: Model identifier (e.g., "gemini-2.0-flash-exp")
+            api_key: Google AI API key (for AI Studio)
+            project_id: GCP project ID (for Vertex AI)
+            location: GCP location (for Vertex AI)
+            use_vertex: Whether to use Vertex AI instead of AI Studio
+
+        Returns:
+            Configured Gemini model instance
+        """
+        return GeminiModel(
+            model_name=model_name,
+            api_key=api_key,
+            project_id=project_id,
+            location=location,
+            use_vertex=use_vertex,
+        )
+
     @staticmethod
     def create_claude_vertex_model(
         model_name: str,
@@ -118,14 +149,57 @@ class ModelFactory:
             credentials=credentials,
         )
 
+    @staticmethod
+    def create_donkit_model(
+        model_name: str,
+        api_key: str,
+        base_url: str = "http://localhost:9017",
+        provider: str = "default",
+    ) -> DonkitModel:
+        """Create a Donkit model that proxies through RagOps API Gateway.
+
+        Args:
+            model_name: Name of the model to use
+            api_key: API key for authentication
+            base_url: Base URL of the RagOps API Gateway
+            provider: Provider to use (default: vertex)
+
+        Returns:
+            DonkitModel instance
+        """
+        return DonkitModel(
+            base_url=base_url,
+            api_token=api_key,
+            provider=provider,
+            model_name=model_name,
+        )
+
     @staticmethod
     def create_model(
         provider: Literal[
-            "openai", "azure_openai", "claude", "claude_vertex", "vertex", "ollama"
+            "openai",
+            "azure_openai",
+            "claude",
+            "claude_vertex",
+            "vertex",
+            "ollama",
+            "donkit",
         ],
-        model_name: str,
+        model_name: str | None,
         credentials: dict,
     ) -> LLMModelAbstract:
+        if model_name is None:
+            default_models = {
+                "openai": "gpt-5-mini",
+                "azure_openai": "gpt-4.1-mini",
+                "claude": "claude-4-5-sonnet",
+                "claude_vertex": "claude-4-5-sonnet",
+                "gemini": "gemini-2.5-flash",
+                "vertex": "gemini-2.5-flash",
+                "ollama": "mistral",
+                "donkit": "default",
+            }
+            model_name = default_models.get(provider, "default")
         if provider == "openai":
             return ModelFactory.create_openai_model(
                 model_name=model_name,
@@ -136,11 +210,19 @@ class ModelFactory:
         elif provider == "azure_openai":
             return ModelFactory.create_azure_openai_model(
                 model_name=model_name,
-                api_key=credentials["api_key"],
-                azure_endpoint=credentials["azure_endpoint"],
+                api_key=credentials.get("api_key"),
+                azure_endpoint=credentials.get("azure_endpoint"),
                 api_version=credentials.get("api_version", "2024-08-01-preview"),
                 deployment_name=credentials.get("deployment_name"),
             )
+        elif provider == "gemini":
+            return ModelFactory.create_gemini_model(
+                model_name=model_name,
+                api_key=credentials.get("api_key"),
+                project_id=credentials.get("project_id"),
+                location=credentials.get("location", "us-central1"),
+                use_vertex=credentials.get("use_vertex", False),
+            )
         elif provider == "claude":
             return ModelFactory.create_claude_model(
                 model_name=model_name,
@@ -162,10 +244,19 @@ class ModelFactory:
             )
         elif provider == "ollama":
             # Ollama uses OpenAI-compatible API
+            ollama_url = credentials.get("ollama_url")
+            if "/v1" not in ollama_url:
+                ollama_url += "/v1"
             return ModelFactory.create_openai_model(
                 model_name=model_name,
                 api_key=credentials.get("api_key", "ollama"),
-                base_url=credentials.get("base_url"),
+                base_url=ollama_url,
+            )
+        elif provider == "donkit":
+            return ModelFactory.create_donkit_model(
+                model_name=model_name,
+                api_key=credentials["api_key"],
+                base_url=credentials["base_url"],
             )
         else:
             raise ValueError(f"Unknown provider: {provider}")