crewplus 0.1.1__tar.gz → 0.1.3__tar.gz

{crewplus-0.1.1 → crewplus-0.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: crewplus
-Version: 0.1.1
+Version: 0.1.3
 Summary: Base services for CrewPlus AI applications
 Author-Email: Tim Liu <tim@opsmateai.com>
 License: MIT
@@ -12,6 +12,9 @@ Requires-Python: <4.0,>=3.11
 Requires-Dist: langchain==0.3.25
 Requires-Dist: langchain-openai==0.3.24
 Requires-Dist: google-genai==1.21.1
+Requires-Dist: mkdocs<2.0.0,>=1.6.1
+Requires-Dist: mkdocs-material<10.0.0,>=9.6.14
+Requires-Dist: mkdocstrings-python<2.0.0,>=1.16.12
 Description-Content-Type: text/markdown
 
 # CrewPlus
@@ -44,6 +47,12 @@ CrewPlus is designed as a modular and extensible ecosystem of packages. This all
 -   **Vector DB Services:** Abstractions for working with popular vector stores for retrieval-augmented generation (RAG).
 -   **Centralized Configuration:** Manage application settings and secrets from a single source of truth (`core/config.py`).
 
+## Documentation
+
+For detailed guides and API references, please see the `docs/` folder.
+
+-   **[GeminiChatModel Documentation](./docs/GeminiChatModel.md)**: A comprehensive guide to using the `GeminiChatModel` for text, image, and video understanding.
+
 ## Installation
 
 To install the core `crewplus` package, run the following command:
@@ -85,9 +94,6 @@ crewplus-base/                    # GitHub repo name
 │       └──  __init__.py
 │       └──  gemini_chat_model.py
 │       └──  model_load_balancer.py
-│       └──  vdb_service.py
-│       └──  ...
-│   └──  vectorstores/
 │       └──  ...
 │   └──  core/
 │       └──  __init__.py
@@ -95,6 +101,8 @@ crewplus-base/                    # GitHub repo name
 │       └──  ...
 ├── tests/
 │   └── ...
+├── docs/
+│   └── ...
 └── notebooks/
     └── ...
 

{crewplus-0.1.1 → crewplus-0.1.3}/README.md

@@ -28,6 +28,12 @@ CrewPlus is designed as a modular and extensible ecosystem of packages. This all
 -   **Vector DB Services:** Abstractions for working with popular vector stores for retrieval-augmented generation (RAG).
 -   **Centralized Configuration:** Manage application settings and secrets from a single source of truth (`core/config.py`).
 
+## Documentation
+
+For detailed guides and API references, please see the `docs/` folder.
+
+-   **[GeminiChatModel Documentation](./docs/GeminiChatModel.md)**: A comprehensive guide to using the `GeminiChatModel` for text, image, and video understanding.
+
 ## Installation
 
 To install the core `crewplus` package, run the following command:
@@ -69,9 +75,6 @@ crewplus-base/                    # GitHub repo name
 │       └──  __init__.py
 │       └──  gemini_chat_model.py
 │       └──  model_load_balancer.py
-│       └──  vdb_service.py
-│       └──  ...
-│   └──  vectorstores/
 │       └──  ...
 │   └──  core/
 │       └──  __init__.py
@@ -79,6 +82,8 @@ crewplus-base/                    # GitHub repo name
 │       └──  ...
 ├── tests/
 │   └── ...
+├── docs/
+│   └── ...
 └── notebooks/
     └── ...
 

crewplus-0.1.3/crewplus/services/gemini_chat_model.py

@@ -0,0 +1,590 @@
+import os
+import asyncio
+import logging
+from typing import Any, Dict, Iterator, List, Optional, AsyncIterator, Union, Tuple
+from google import genai
+from google.genai import types
+import base64
+import requests
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import (
+    AIMessage,
+    AIMessageChunk,
+    BaseMessage,
+    HumanMessage,
+    SystemMessage,
+)
+from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
+from langchain_core.callbacks import (
+    CallbackManagerForLLMRun,
+    AsyncCallbackManagerForLLMRun
+)
+from pydantic import Field, SecretStr
+from langchain_core.utils import convert_to_secret_str
+
+class GeminiChatModel(BaseChatModel):
+    """Custom chat model for Google Gemini, supporting text, image, and video.
+    
+    This model provides a robust interface to Google's Gemini Pro and Flash models,
+    handling various data formats for multimodal inputs while maintaining compatibility
+    with the LangChain ecosystem.
+
+    It supports standard invocation, streaming, and asynchronous operations.
+    API keys can be provided directly or loaded from the `GOOGLE_API_KEY`
+    environment variable.
+
+    Attributes:
+        model_name (str): The Google model name to use (e.g., "gemini-1.5-flash").
+        google_api_key (Optional[SecretStr]): Your Google API key.
+        temperature (Optional[float]): The sampling temperature for generation.
+        max_tokens (Optional[int]): The maximum number of tokens to generate.
+        top_p (Optional[float]): The top-p (nucleus) sampling parameter.
+        top_k (Optional[int]): The top-k sampling parameter.
+        logger (Optional[logging.Logger]): An optional logger instance.
+
+    Example:
+        .. code-block:: python
+
+            from crewplus.services import GeminiChatModel
+            from langchain_core.messages import HumanMessage
+            import base64
+            import logging
+
+            # Initialize the model with optional logger
+            logger = logging.getLogger("my_app.gemini")
+            model = GeminiChatModel(model_name="gemini-2.0-flash", logger=logger)
+
+            # --- Text-only usage ---
+            response = model.invoke("Hello, how are you?")
+            print("Text response:", response.content)
+
+            # --- Image processing with base64 data URI ---
+            # Replace with a path to your image
+            image_path = "path/to/your/image.jpg"
+            try:
+                with open(image_path, "rb") as image_file:
+                    encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
+                
+                image_message = HumanMessage(
+                    content=[
+                        {"type": "text", "text": "What is in this image?"},
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/jpeg;base64,{encoded_string}"
+                            }
+                        },
+                    ]
+                )
+                image_response = model.invoke([image_message])
+                print("Image response (base64):", image_response.content)
+            except FileNotFoundError:
+                print(f"Image file not found at {image_path}, skipping base64 example.")
+
+
+            # --- Image processing with URL ---
+            url_message = HumanMessage(
+                content=[
+                    {"type": "text", "text": "Describe this image:"},
+                    {
+                        "type": "image_url",
+                        "image_url": "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
+                    },
+                ]
+            )
+            url_response = model.invoke([url_message])
+            print("Image response (URL):", url_response.content)
+            
+            # --- Video processing with file path (>=20MB) ---
+            video_path = "path/to/your/video.mp4"
+            video_file = client.files.upload(file=video_path)
+
+            try:
+                video_message = HumanMessage(
+                    content=[
+                        {"type": "text", "text": "Summarize this video."},
+                        {"type": "video_file", "file": video_file},
+                    ]
+                )
+                video_response = model.invoke([video_message])
+                print("Video response (file path):", video_response.content)
+            except Exception as e:
+                print(f"Video processing with file path failed: {e}")
+
+            # --- Video processing with raw bytes (<20MB) ---
+            video_path = "path/to/your/video.mp4"
+            try:
+                with open(video_path, "rb") as video_file:
+                    video_bytes = video_file.read()
+                
+                video_message = HumanMessage(
+                    content=[
+                        {"type": "text", "text": "What is happening in this video?"},
+                        {
+                            "type": "video_file",
+                            "data": video_bytes,
+                            "mime_type": "video/mp4"
+                        },
+                    ]
+                )
+                video_response = model.invoke([video_message])
+                print("Video response (bytes):", video_response.content)
+            except FileNotFoundError:
+                print(f"Video file not found at {video_path}, skipping bytes example.")
+            except Exception as e:
+                print(f"Video processing with bytes failed: {e}")
+            
+            # --- Streaming usage (works with text, images, and video) ---
+            print("Streaming response:")
+            for chunk in model.stream([url_message]):
+                print(chunk.content, end="", flush=True)
+    """
+    
+    # Model configuration
+    model_name: str = Field(default="gemini-2.0-flash", description="The Google model name to use")
+    google_api_key: Optional[SecretStr] = Field(default=None, description="Google API key")
+    temperature: Optional[float] = Field(default=0.7, description="Sampling temperature")
+    max_tokens: Optional[int] = Field(default=None, description="Maximum tokens to generate")
+    top_p: Optional[float] = Field(default=None, description="Top-p sampling parameter")
+    top_k: Optional[int] = Field(default=None, description="Top-k sampling parameter")
+    logger: Optional[logging.Logger] = Field(default=None, description="Optional logger instance")
+    
+    # Internal client
+    _client: Optional[genai.Client] = None
+    
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        
+        # Initialize logger
+        if self.logger is None:
+            self.logger = logging.getLogger(f"{self.__class__.__module__}.{self.__class__.__name__}")
+            if not self.logger.handlers: # and not getattr(self.logger, 'propagate', True):
+                self.logger.addHandler(logging.StreamHandler())
+                self.logger.setLevel(logging.INFO)
+        
+        # Get API key from environment if not provided
+        if self.google_api_key is None:
+            api_key = os.getenv("GOOGLE_API_KEY")
+            if api_key:
+                self.google_api_key = convert_to_secret_str(api_key)
+        
+        # Initialize the Google GenAI client
+        if self.google_api_key:
+            self._client = genai.Client(
+                api_key=self.google_api_key.get_secret_value()
+            )
+            self.logger.info(f"Initialized GeminiChatModel with model: {self.model_name}")
+        else:
+            error_msg = "Google API key is required. Set GOOGLE_API_KEY environment variable or pass google_api_key parameter."
+            self.logger.error(error_msg)
+            raise ValueError(error_msg)
+    
+    @property
+    def _llm_type(self) -> str:
+        """Return identifier for the model type."""
+        return "custom_google_genai"
+    
+    @property
+    def _identifying_params(self) -> Dict[str, Any]:
+        """Return a dictionary of identifying parameters for tracing."""
+        return {
+            "model_name": self.model_name,
+            "temperature": self.temperature,
+            "max_tokens": self.max_tokens,
+            "top_p": self.top_p,
+            "top_k": self.top_k,
+        }
+    
+    def _convert_messages(self, messages: List[BaseMessage]) -> Union[types.ContentListUnion, types.ContentListUnionDict]:
+        """
+        Converts LangChain messages to a format suitable for the GenAI API.
+        - For single, multi-part HumanMessage, returns a direct list of parts (e.g., [File, "text"]).
+        - For multi-turn chats, returns a list of Content objects.
+        - For simple text, returns a string.
+        """
+        self.logger.debug(f"Converting {len(messages)} messages.")
+        
+        # Filter out system messages (handled in generation_config)
+        chat_messages = [msg for msg in messages if not isinstance(msg, SystemMessage)]
+
+        # Case 1: A single HumanMessage. This is the most common path for single prompts.
+        if len(chat_messages) == 1 and isinstance(chat_messages[0], HumanMessage):
+            content = chat_messages[0].content
+            # For a simple string, return it directly.
+            if isinstance(content, str):
+                return content
+            # For a list of parts, parse them into a direct list for the API.
+            return list(self._parse_message_content(content, is_simple=True))
+
+
+        # Case 2: Multi-turn chat history. This requires a list of Content objects.
+        self.logger.debug("Handling as a multi-turn chat conversation.")
+        genai_contents: List[types.Content] = []
+        for msg in chat_messages:
+            role = "model" if isinstance(msg, AIMessage) else "user"
+            parts = []
+            
+            # Process each part and ensure proper typing
+            for part in self._parse_message_content(msg.content, is_simple=False):
+                if isinstance(part, types.File):
+                    # put File directly into types.Content
+                    parts.append(part)
+                elif isinstance(part, types.Part):
+                    parts.append(part)
+                else:
+                    self.logger.warning(f"Unexpected part type: {type(part)}")
+            
+            if parts:
+                genai_contents.append(types.Content(parts=parts, role=role))
+        
+        # If there's only one Content object, return it directly instead of a list
+        if len(genai_contents) == 1:
+            return genai_contents[0]
+            
+        return genai_contents
+
+    def _create_image_part(self, image_info: Dict[str, Any]) -> Union[types.Part, types.File]:
+        """Creates a GenAI Part or File from various image source formats."""
+        self.logger.debug(f"Creating image part from info: {list(image_info.keys())}")
+
+        if "path" in image_info:
+            return self._client.files.upload(file=image_info["path"])
+        
+        if "data" in image_info:
+            data = image_info["data"]
+            if image_info.get("source_type") == "base64":
+                data = base64.b64decode(data)
+            return types.Part.from_bytes(data=data, mime_type=image_info["mime_type"])
+
+        url = image_info.get("image_url", image_info.get("url"))
+        if isinstance(url, dict):
+            url = url.get("url")
+        
+        if not url:
+            raise ValueError(f"Invalid image info, requires 'path', 'data', or 'url'. Received: {image_info}")
+
+        if url.startswith("data:"):
+            header, encoded = url.split(",", 1)
+            mime_type = header.split(":", 1)[-1].split(";", 1)[0]
+            image_data = base64.b64decode(encoded)
+            return types.Part.from_bytes(data=image_data, mime_type=mime_type)
+        else:
+            response = requests.get(url)
+            response.raise_for_status()
+            mime_type = response.headers.get("Content-Type", "image/jpeg")
+            return types.Part.from_bytes(data=response.content, mime_type=mime_type)
+
+    def _create_video_part(self, video_info: Dict[str, Any]) -> Union[types.Part, types.File]:
+        """Creates a Google GenAI Part or File from video information.
+        
+        Supports multiple video input formats:
+        - File object: {"type": "video_file", "file": file_object}
+        - File path: {"type": "video_file", "path": "/path/to/video.mp4"}
+        - Raw bytes: {"type": "video_file", "data": video_bytes, "mime_type": "video/mp4"}
+        - URL/URI: {"type": "video_file", "url": "https://example.com/video.mp4"}
+        - YouTube URL: {"type": "video_file", "url": "https://www.youtube.com/watch?v=..."}
+        - URL with offset: {"type": "video_file", "url": "...", "start_offset": "12s", "end_offset": "50s"}
+        
+        Args:
+            video_info: Dictionary containing video information
+            
+        Returns:
+            Either a types.Part or File object for Google GenAI
+            
+        Raises:
+            FileNotFoundError: If video file path doesn't exist
+            ValueError: If video_info is invalid or missing required fields
+        """
+        self.logger.debug(f"Creating video part from info: {list(video_info.keys())}")
+        
+        # Handle pre-uploaded file object
+        if "file" in video_info:
+            if isinstance(video_info["file"], types.File):
+                return video_info["file"]
+            else:
+                raise ValueError(f"The 'file' key must contain a google.genai.File object, but got {type(video_info['file'])}")
+
+        if "path" in video_info:
+            self.logger.debug(f"Uploading video file from path: {video_info['path']}")
+
+            uploaded_file =self._client.files.upload(file=video_info["path"])
+
+            self.logger.debug(f"Uploaded video file: {uploaded_file}")
+
+            return uploaded_file
+        
+        mime_type = video_info.get("mime_type")
+
+        if "data" in video_info:
+            data = video_info["data"]
+            if not mime_type:
+                raise ValueError("'mime_type' is required when providing video data.")
+            max_size = 20 * 1024 * 1024  # 20MB
+            if len(data) > max_size:
+                raise ValueError(f"Video data size ({len(data)} bytes) exceeds 20MB limit for inline data.")
+            return types.Part(inline_data=types.Blob(data=data, mime_type=mime_type))
+
+        url = video_info.get("url")
+        if not url:
+            raise ValueError(f"Invalid video info, requires 'path', 'data', 'url', or 'file'. Received: {video_info}")
+
+        mime_type = video_info.get("mime_type", "video/mp4")
+        
+        # Handle video offsets
+        start_offset = video_info.get("start_offset")
+        end_offset = video_info.get("end_offset")
+
+        self.logger.debug(f"Video offsets: {start_offset} to {end_offset}.")
+        
+        if start_offset or end_offset:
+            video_metadata = types.VideoMetadata(start_offset=start_offset, end_offset=end_offset)
+            return types.Part(
+                file_data=types.FileData(file_uri=url, mime_type=mime_type),
+                video_metadata=video_metadata
+            )
+
+        return types.Part(file_data=types.FileData(file_uri=url, mime_type=mime_type))
+
+    def _parse_message_content(
+        self, content: Union[str, List[Union[str, Dict]]], *, is_simple: bool = True
+    ) -> Iterator[Union[str, types.Part, types.File]]:
+        """
+        Parses LangChain message content and yields parts for Google GenAI.
+
+        Args:
+            content: The message content to parse.
+            is_simple: If True, yields raw objects where possible (e.g., str, File)
+                               for single-turn efficiency. If False, ensures all yielded
+                               parts are `types.Part` by converting raw strings and
+                               Files as needed, which is required for multi-turn chat.
+
+        Supports both standard LangChain formats and enhanced video formats:
+        - Text: "string" or {"type": "text", "text": "content"}
+        - Image: {"type": "image_url", "image_url": "url"} or {"type": "image_url", "image_url": {"url": "url"}}
+        - Video: {"type": "video_file", ...} or {"type": "video", ...}                               
+        """
+        if isinstance(content, str):
+            yield content if is_simple else types.Part(text=content)
+            return
+
+        if not isinstance(content, list):
+            self.logger.warning(f"Unsupported content format: {type(content)}")
+            return
+
+        for i, part_spec in enumerate(content):
+            try:
+                if isinstance(part_spec, str):
+                    yield part_spec if is_simple else types.Part(text=part_spec)
+                    continue
+                
+                if isinstance(part_spec, types.File):
+                    if is_simple:
+                        yield part_spec
+                    else:
+                        yield types.Part(file_data=types.FileData(
+                            mime_type=part_spec.mime_type,
+                            file_uri=part_spec.uri
+                        ))
+                    continue
+
+                if not isinstance(part_spec, dict):
+                    self.logger.warning(f"Skipping non-dict part in content list: {type(part_spec)}")
+                    continue
+
+                part_type = part_spec.get("type", "").lower()
+                
+                if part_type == "text":
+                    if text_content := part_spec.get("text"):
+                        yield text_content if is_simple else types.Part(text=text_content)
+                elif part_type in ("image", "image_url"):
+                    yield self._create_image_part(part_spec)
+                elif part_type in ("video", "video_file"):
+                    yield self._create_video_part(part_spec)
+                else:
+                    self.logger.debug(f"Part with unknown type '{part_type}' was ignored at index {i}.")
+            except Exception as e:
+                self.logger.error(f"Failed to process message part at index {i}: {part_spec}. Error: {e}", exc_info=True)
+
+    def _prepare_generation_config(
+        self, messages: List[BaseMessage], stop: Optional[List[str]] = None
+    ) -> Dict[str, Any]:
+        """Prepares the generation configuration, including system instructions."""
+        # Base config from model parameters
+        config = {
+            "temperature": self.temperature,
+            "max_output_tokens": self.max_tokens,
+            "top_p": self.top_p,
+            "top_k": self.top_k,
+        }
+        if stop:
+            config["stop_sequences"] = stop
+        
+        # Handle system instructions
+        system_prompts = [msg.content for msg in messages if isinstance(msg, SystemMessage) and msg.content]
+        if system_prompts:
+            system_prompt_str = "\n\n".join(system_prompts)
+            config["system_instruction"] = system_prompt_str
+        
+        # Filter out None values before returning
+        return {k: v for k, v in config.items() if v is not None}
+
+    def _trim_for_logging(self, contents: Any) -> Any:
+        """Helper to trim large binary data from logging payloads."""
+        if isinstance(contents, str):
+            return contents
+        
+        if isinstance(contents, types.Content):
+            return {
+                "role": contents.role,
+                "parts": [self._trim_part(part) for part in contents.parts]
+            }
+        
+        if isinstance(contents, list):
+            return [self._trim_for_logging(item) for item in contents]
+        
+        return contents
+
+    def _trim_part(self, part: types.Part) -> dict:
+        """Trims individual part data for safe logging."""
+        part_dict = {}
+        if part.text:
+            part_dict["text"] = part.text
+        if part.inline_data:
+            part_dict["inline_data"] = {
+                "mime_type": part.inline_data.mime_type,
+                "data_size": f"{len(part.inline_data.data)} bytes"
+            }
+        if part.file_data:
+            part_dict["file_data"] = {
+                "mime_type": part.file_data.mime_type,
+                "file_uri": part.file_data.file_uri
+            }
+        return part_dict
+
+    def _generate(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        """Generates a chat response from a list of messages."""
+        self.logger.info(f"Generating response for {len(messages)} messages.")
+        
+        contents = self._convert_messages(messages)
+        config = self._prepare_generation_config(messages, stop)
+
+        try:
+            response = self._client.models.generate_content(
+                model=self.model_name,
+                contents=contents,
+                config=config,
+                **kwargs,
+            )
+            
+            generated_text = response.text
+            finish_reason = response.candidates[0].finish_reason.name if response.candidates else None
+            
+            message = AIMessage(
+                content=generated_text,
+                response_metadata={"model_name": self.model_name, "finish_reason": finish_reason},
+            )
+            return ChatResult(generations=[ChatGeneration(message=message)])
+            
+        except Exception as e:
+            self.logger.error(f"Error generating content with Google GenAI: {e}", exc_info=True)
+            raise ValueError(f"Error during generation: {e}")
+
+    async def _agenerate(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        """Asynchronously generates a chat response."""
+        self.logger.info(f"Async generating response for {len(messages)} messages.")
+        
+        contents = self._convert_messages(messages)
+        config = self._prepare_generation_config(messages, stop)
+
+        try:
+            response = await self._client.generate_content(
+                model=self.model_name,
+                contents=contents,
+                config=config,
+                **kwargs,
+            )
+            
+            generated_text = response.text
+            finish_reason = response.candidates[0].finish_reason.name if response.candidates else None
+            
+            message = AIMessage(
+                content=generated_text,
+                response_metadata={"model_name": self.model_name, "finish_reason": finish_reason},
+            )
+            return ChatResult(generations=[ChatGeneration(message=message)])
+
+        except Exception as e:
+            self.logger.error(f"Error during async generation: {e}", exc_info=True)
+            raise ValueError(f"Error during async generation: {e}")
+
+    def _stream(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> Iterator[ChatGenerationChunk]:
+        """Streams the chat response."""
+        self.logger.info(f"Streaming response for {len(messages)} messages.")
+        
+        contents = self._convert_messages(messages)
+        config = self._prepare_generation_config(messages, stop)
+
+        try:
+            stream = self._client.models.generate_content_stream(
+                model=self.model_name,
+                contents=contents,
+                config=config,
+                **kwargs,
+            )
+            for chunk_response in stream:
+                if text_content := chunk_response.text:
+                    chunk = ChatGenerationChunk(message=AIMessageChunk(content=text_content))
+                    if run_manager:
+                        run_manager.on_llm_new_token(text_content, chunk=chunk)
+                    yield chunk
+        except Exception as e:
+            self.logger.error(f"Error streaming content: {e}", exc_info=True)
+            raise ValueError(f"Error during streaming: {e}")
+
+    async def _astream(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[ChatGenerationChunk]:
+        """Asynchronously streams the chat response."""
+        self.logger.info(f"Async streaming response for {len(messages)} messages.")
+        
+        contents = self._convert_messages(messages)
+        config = self._prepare_generation_config(messages, stop)
+
+        try:
+            stream = await self._client.generate_content_async(
+                model=self.model_name,
+                contents=contents,
+                config=config,
+                **kwargs,
+            )
+            async for chunk_response in stream:
+                if text_content := chunk_response.text:
+                    chunk = ChatGenerationChunk(message=AIMessageChunk(content=text_content))
+                    if run_manager:
+                        await run_manager.on_llm_new_token(text_content, chunk=chunk)
+                    yield chunk
+        except Exception as e:
+            self.logger.error(f"Error during async streaming: {e}", exc_info=True)
+            raise ValueError(f"Error during async streaming: {e}")

crewplus-0.1.3/crewplus/services/model_load_balancer.py

@@ -0,0 +1,184 @@
+import json
+import random
+import logging
+from typing import Dict, List, Optional, Union
+from collections import defaultdict
+from langchain_openai import AzureChatOpenAI, ChatOpenAI, AzureOpenAIEmbeddings
+from .gemini_chat_model import GeminiChatModel
+
+
+class ModelLoadBalancer:
+    def __init__(self,
+                 config_path: Optional[str] = "config/models_config.json",
+                 config_data: Optional[Dict] = None,
+                 logger: Optional[logging.Logger] = None):
+        """
+        Initializes the ModelLoadBalancer.
+
+        Args:
+            config_path: Path to the JSON configuration file.
+            config_data: A dictionary containing the model configuration.
+            logger: An optional logger instance. If not provided, a default one is created.
+
+        Raises:
+            ValueError: If neither config_path nor config_data is provided.
+        """
+        if not config_path and not config_data:
+            raise ValueError("Either 'config_path' or 'config_data' must be provided.")
+
+        self.config_path = config_path
+        self.config_data = config_data
+        self.logger = logger or logging.getLogger(__name__)
+        self.models_config: List[Dict] = []
+        self.models: Dict[int, Union[AzureChatOpenAI, ChatOpenAI, AzureOpenAIEmbeddings, GeminiChatModel]] = {}
+        self._initialize_state()
+        self._config_loaded = False  # Flag to check if config is loaded
+
+    def load_config(self):
+        """Load and validate model configurations from a file path or a dictionary."""
+        self.logger.debug("Model balancer: loading configuration.")
+        try:
+            config = None
+            if self.config_data:
+                config = self.config_data
+            elif self.config_path:
+                with open(self.config_path, 'r') as f:
+                    config = json.load(f)
+            else:
+                # This case is handled in __init__, but as a safeguard:
+                raise RuntimeError("No configuration source provided (path or data).")
+
+            # Validate config
+            if 'models' not in config or not isinstance(config['models'], list):
+                raise ValueError("Configuration must contain a 'models' list.")
+
+            for model in config.get('models', []):
+                if 'provider' not in model or 'type' not in model or 'id' not in model:
+                    self.logger.error("Model config must contain 'id', 'provider', and 'type' fields.")
+                    raise ValueError("Model config must contain 'id', 'provider', and 'type' fields.")
+
+            self.models_config = config['models']
+
+            # Instantiate models
+            for model_config in self.models_config:
+                model_id = model_config['id']
+                self.models[model_id] = self._instantiate_model(model_config)
+
+            self._config_loaded = True
+            self.logger.debug("Model balancer: configuration loaded successfully.")
+        except (FileNotFoundError, json.JSONDecodeError, ValueError) as e:
+            self._config_loaded = False
+            self.logger.error(f"Failed to load model configuration: {e}", exc_info=True)
+            raise RuntimeError(f"Failed to load model configuration: {e}")
+
+    def get_model(self, provider: str = None, model_type: str = None, deployment_name: str = None):
+        """
+        Get a model instance.
+
+        Can fetch a model in two ways:
+        1. By its specific `deployment_name`.
+        2. By `provider` and `model_type`, which will select a model using round-robin.
+
+        Args:
+            provider: The model provider (e.g., 'azure-openai', 'google-genai').
+            model_type: The type of model (e.g., 'inference', 'embedding').
+            deployment_name: The unique name for the model deployment.
+
+        Returns:
+            An instantiated language model object.
+
+        Raises:
+            RuntimeError: If the model configuration has not been loaded.
+            ValueError: If the requested model cannot be found or if parameters are insufficient.
+        """
+        if not self._config_loaded:
+            self.logger.error("Model configuration not loaded")
+            raise RuntimeError("Model configuration not loaded")
+
+        if deployment_name:
+            for model_config in self.models_config:
+                if model_config.get('deployment_name') == deployment_name:
+                    model_id = model_config['id']
+                    return self.models[model_id]
+            self.logger.error(f"No model found for deployment name: {deployment_name}")
+            raise ValueError(f"No model found for deployment name: {deployment_name}")
+
+        if provider and model_type:
+            candidates = [model for model in self.models_config if model.get('provider') == provider and model.get('type') == model_type]
+            if not candidates:
+                self.logger.error(f"No models found for provider '{provider}' and type '{model_type}'")
+                raise ValueError(f"No models found for provider '{provider}' and type '{model_type}'")
+
+            selected_model_config = self._round_robin_selection(candidates)
+            model_id = selected_model_config['id']
+            return self.models[model_id]
+
+        raise ValueError("Either 'deployment_name' or both 'provider' and 'model_type' must be provided.")
+
+    def _instantiate_model(self, model_config: Dict):
+        """Instantiate and return an LLM object based on the model configuration"""
+        provider = model_config['provider']
+        self.logger.debug(f"Model balancer: instantiating {provider} -- {model_config.get('deployment_name')}")
+
+        if provider == 'azure-openai':
+            kwargs = {
+                'azure_deployment': model_config['deployment_name'],
+                'openai_api_version': model_config['api_version'],
+                'azure_endpoint': model_config['api_base'],
+                'openai_api_key': model_config['api_key']
+            }
+            if 'temperature' in model_config:
+                kwargs['temperature'] = model_config['temperature']
+            if model_config.get('deployment_name') == 'o1-mini':
+                kwargs['disable_streaming'] = True
+            return AzureChatOpenAI(**kwargs)
+        elif provider == 'openai':
+            kwargs = {
+                'openai_api_key': model_config['api_key']
+            }
+            if 'temperature' in model_config:
+                kwargs['temperature'] = model_config['temperature']
+            return ChatOpenAI(**kwargs)
+        elif provider == 'azure-openai-embeddings':
+            return AzureOpenAIEmbeddings(
+                azure_deployment=model_config['deployment_name'],
+                openai_api_version=model_config['api_version'],
+                api_key=model_config['api_key'],
+                azure_endpoint=model_config['api_base'],
+                chunk_size=16, request_timeout=60, max_retries=2
+            )
+        elif provider == 'google-genai':
+            kwargs = {
+                'google_api_key': model_config['api_key'],
+                'model_name': model_config['deployment_name']  # Map deployment_name to model_name
+            }
+            if 'temperature' in model_config:
+                kwargs['temperature'] = model_config['temperature']
+            if 'max_tokens' in model_config:
+                kwargs['max_tokens'] = model_config['max_tokens']
+            return GeminiChatModel(**kwargs)
+        else:
+            self.logger.error(f"Unsupported provider: {provider}")
+            raise ValueError(f"Unsupported provider: {provider}")
+
+    def _initialize_state(self):
+        self.active_models = []
+        self.usage_counter = defaultdict(int)
+        self.current_indices = {}
+
+    def _round_robin_selection(self, candidates: list) -> Dict:
+        if id(candidates) not in self.current_indices:
+            self.current_indices[id(candidates)] = 0
+        idx = self.current_indices[id(candidates)]
+        model = candidates[idx]
+        self.current_indices[id(candidates)] = (idx + 1) % len(candidates)
+        self.usage_counter[model['id']] += 1
+
+        return model
+
+    def _least_used_selection(self, candidates: list) -> Dict:
+        min_usage = min(self.usage_counter[m['model_id']] for m in candidates)
+        least_used = [m for m in candidates if self.usage_counter[m['model_id']] == min_usage]
+        model = random.choice(least_used)
+        self.usage_counter[model['id']] += 1
+        return model

{crewplus-0.1.1 → crewplus-0.1.3}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "pdm.backend"
 
 [project]
 name = "crewplus"
-version = "0.1.1"
+version = "0.1.3"
 description = "Base services for CrewPlus AI applications"
 authors = [
     { name = "Tim Liu", email = "tim@opsmateai.com" },
@@ -17,6 +17,9 @@ dependencies = [
     "langchain==0.3.25",
     "langchain-openai==0.3.24",
     "google-genai==1.21.1",
+    "mkdocs (>=1.6.1,<2.0.0)",
+    "mkdocs-material (>=9.6.14,<10.0.0)",
+    "mkdocstrings-python (>=1.16.12,<2.0.0)",
 ]
 
 [project.license]

crewplus-0.1.3/tests/models_config.json

@@ -0,0 +1,77 @@
+{
+    "models": [
+        {
+            "id": 1,
+            "provider": "azure-openai",
+            "type": "inference",
+            "deployment_name": "gpt-o3mini-eastus2-RPM25",
+            "api_version": "2024-12-01-preview",
+            "api_base": "https://crewplus-eastus2.openai.azure.com",
+            "api_key": "c67cb5d0d8ae4aef81d7f42aeae274b6"
+        },
+        {
+            "id": 2,
+            "provider": "azure-openai",
+            "type": "ingestion",
+            "deployment_name": "gpt-4o",
+            "api_version": "2025-01-01-preview",
+            "api_base": "https://crewplus-eastus2.openai.azure.com",
+            "api_key": "c67cb5d0d8ae4aef81d7f42aeae274b6",
+            "temperature": 0.0
+        },    
+        {
+            "id": 3,
+            "provider": "azure-openai",
+            "type": "inference",
+            "deployment_name": "gpt-4.1",
+            "api_version": "2025-01-01-preview",
+            "api_base": "https://crewplus-eastus2.openai.azure.com",
+            "api_key": "c67cb5d0d8ae4aef81d7f42aeae274b6",
+            "temperature": 0.0
+        },             
+        {
+            "id": 4,
+            "provider": "azure-openai",
+            "type": "ingestion",
+            "deployment_name": "cpai-gpt4o-westus",
+            "api_version": "2025-01-01-preview",
+            "api_base": "https://crewplus-westus.openai.azure.com",
+            "api_key": "b93bc4d2ef8e4298bd8390002922d084",
+            "temperature": 0.0
+        },   
+        {
+            "id": 5,
+            "provider": "azure-openai-embeddings",
+            "type": "embedding",
+            "deployment_name": "cpai-text-embedding-ada-002-westus",
+            "api_version": "2024-02-01",
+            "api_base": "https://crewplus-westus.openai.azure.com",
+            "api_key": "b93bc4d2ef8e4298bd8390002922d084"
+        },       
+        {
+            "id": 6,
+            "provider": "azure-openai-embeddings",
+            "type": "embedding",
+            "deployment_name": "cpai-text-embedding-3-large-eastus2",
+            "api_version": "1",
+            "api_base": "https://crewplus-eastus2.openai.azure.com",
+            "api_key": "c67cb5d0d8ae4aef81d7f42aeae274b6"
+        },
+        {
+            "id": 7,
+            "provider": "google-genai",
+            "type": "inference",
+            "deployment_name": "gemini-2.5-flash",
+            "api_key": "AIzaSyDkZbcGcV7SB6OyN4XkK_sF2mzO2E-nKQk",
+            "temperature": 0.0
+        },
+        {
+            "id": 8,
+            "provider": "google-genai",
+            "type": "ingestion",
+            "deployment_name": "gemini-2.5-pro",
+            "api_key": "AIzaSyDkZbcGcV7SB6OyN4XkK_sF2mzO2E-nKQk",
+            "temperature": 0.0
+        }
+    ]
+}

crewplus-0.1.1/crewplus/services/gemini_chat_model.py

@@ -1,365 +0,0 @@
-import os
-import asyncio
-from typing import Any, Dict, Iterator, List, Optional, AsyncIterator
-from google import genai
-from langchain_core.language_models import BaseChatModel
-from langchain_core.messages import (
-    AIMessage,
-    AIMessageChunk,
-    BaseMessage,
-    HumanMessage,
-    SystemMessage,
-)
-from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
-from langchain_core.callbacks import (
-    CallbackManagerForLLMRun,
-    AsyncCallbackManagerForLLMRun
-)
-from pydantic import Field, SecretStr
-from langchain_core.utils import convert_to_secret_str
-
-class GeminiChatModel(BaseChatModel):
-    """Custom chat model using Google's genai client package directly with real streaming support.
-    
-    This implementation provides direct access to Google's genai features
-    while being compatible with LangChain's BaseChatModel interface.
-    
-    Example:
-        ```python
-        model = GeminiChatModel(
-            model_name="gemini-2.0-flash",
-            google_api_key="your-api-key",
-            temperature=0.7
-        )
-        
-        # Basic usage
-        response = model.invoke("Hello, how are you?")
-        print(response.content)
-        
-        # Streaming usage
-        for chunk in model.stream("Tell me a story"):
-            print(chunk.content, end="")
-            
-        # Async usage
-        async def test_async():
-            response = await model.ainvoke("Hello!")
-            print(response.content)
-            
-            async for chunk in model.astream("Tell me a story"):
-                print(chunk.content, end="")
-        ```
-    """
-    
-    # Model configuration
-    model_name: str = Field(default="gemini-2.0-flash", description="The Google model name to use")
-    google_api_key: Optional[SecretStr] = Field(default=None, description="Google API key")
-    temperature: Optional[float] = Field(default=0.7, description="Sampling temperature")
-    max_tokens: Optional[int] = Field(default=None, description="Maximum tokens to generate")
-    top_p: Optional[float] = Field(default=None, description="Top-p sampling parameter")
-    top_k: Optional[int] = Field(default=None, description="Top-k sampling parameter")
-    
-    # Internal client
-    _client: Optional[genai.Client] = None
-    
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        
-        # Get API key from environment if not provided
-        if self.google_api_key is None:
-            api_key = os.getenv("GOOGLE_API_KEY")
-            if api_key:
-                self.google_api_key = convert_to_secret_str(api_key)
-        
-        # Initialize the Google GenAI client
-        if self.google_api_key:
-            self._client = genai.Client(
-                api_key=self.google_api_key.get_secret_value()
-            )
-        else:
-            raise ValueError("Google API key is required. Set GOOGLE_API_KEY environment variable or pass google_api_key parameter.")
-    
-    @property
-    def _llm_type(self) -> str:
-        """Return identifier for the model type."""
-        return "custom_google_genai"
-    
-    @property
-    def _identifying_params(self) -> Dict[str, Any]:
-        """Return a dictionary of identifying parameters for tracing."""
-        return {
-            "model_name": self.model_name,
-            "temperature": self.temperature,
-            "max_tokens": self.max_tokens,
-            "top_p": self.top_p,
-            "top_k": self.top_k,
-        }
-    
-    def _convert_messages_to_genai_format(self, messages: List[BaseMessage]) -> str:
-        """Convert LangChain messages to Google GenAI format.
-        
-        Google GenAI API doesn't support system messages, so we'll convert
-        the conversation to a single prompt string with proper formatting.
-        """
-        prompt_parts = []
-        
-        for message in messages:
-            if isinstance(message, SystemMessage):
-                # Convert system message to instruction format
-                prompt_parts.append(f"Instructions: {message.content}")
-            elif isinstance(message, HumanMessage):
-                prompt_parts.append(f"Human: {message.content}")
-            elif isinstance(message, AIMessage):
-                prompt_parts.append(f"Assistant: {message.content}")
-            else:
-                # Default to human format for unknown message types
-                prompt_parts.append(f"Human: {str(message.content)}")
-        
-        # Add a final prompt for the assistant to respond
-        if not prompt_parts or not prompt_parts[-1].startswith("Human:"):
-            prompt_parts.append("Human: Please respond to the above.")
-        
-        prompt_parts.append("Assistant:")
-        
-        return "\n\n".join(prompt_parts)
-    
-    def _prepare_generation_config(self, stop: Optional[List[str]] = None) -> Dict[str, Any]:
-        """Prepare generation configuration for Google GenAI."""
-        generation_config = {}
-        if self.temperature is not None:
-            generation_config["temperature"] = self.temperature
-        if self.max_tokens is not None:
-            generation_config["max_output_tokens"] = self.max_tokens
-        if self.top_p is not None:
-            generation_config["top_p"] = self.top_p
-        if self.top_k is not None:
-            generation_config["top_k"] = self.top_k
-        if stop:
-            generation_config["stop_sequences"] = stop
-        return generation_config
-    
-    def _generate(
-        self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> ChatResult:
-        """Generate a response using Google's genai client."""
-        
-        # Convert messages to a single prompt string
-        prompt = self._convert_messages_to_genai_format(messages)
-        
-        # Prepare generation config
-        generation_config = self._prepare_generation_config(stop)
-        
-        try:
-            # Generate response using Google GenAI
-            response = self._client.models.generate_content(
-                model=self.model_name,
-                contents=prompt,
-                config=generation_config if generation_config else None
-            )
-            
-            # Extract the generated text
-            generated_text = response.text if hasattr(response, 'text') else str(response)
-            
-            # Create AI message with response metadata
-            message = AIMessage(
-                content=generated_text,
-                response_metadata={
-                    "model_name": self.model_name,
-                    "finish_reason": getattr(response, 'finish_reason', None),
-                }
-            )
-            
-            # Create and return ChatResult
-            generation = ChatGeneration(message=message)
-            return ChatResult(generations=[generation])
-            
-        except Exception as e:
-            raise ValueError(f"Error generating content with Google GenAI: {str(e)}")
-    
-    async def _agenerate(
-        self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> ChatResult:
-        """Async generate a response using Google's genai client."""
-        
-        # Convert messages to a single prompt string
-        prompt = self._convert_messages_to_genai_format(messages)
-        
-        # Prepare generation config
-        generation_config = self._prepare_generation_config(stop)
-        
-        try:
-            # Generate response using Google GenAI (run in executor for async)
-            loop = asyncio.get_event_loop()
-            response = await loop.run_in_executor(
-                None,
-                lambda: self._client.models.generate_content(
-                    model=self.model_name,
-                    contents=prompt,
-                    config=generation_config if generation_config else None
-                )
-            )
-            
-            # Extract the generated text
-            generated_text = response.text if hasattr(response, 'text') else str(response)
-            
-            # Create AI message with response metadata
-            message = AIMessage(
-                content=generated_text,
-                response_metadata={
-                    "model_name": self.model_name,
-                    "finish_reason": getattr(response, 'finish_reason', None),
-                }
-            )
-            
-            # Create and return ChatResult
-            generation = ChatGeneration(message=message)
-            return ChatResult(generations=[generation])
-            
-        except Exception as e:
-            raise ValueError(f"Error generating content with Google GenAI: {str(e)}")
-    
-    def _stream(
-        self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> Iterator[ChatGenerationChunk]:
-        """Stream the output using Google's genai client with real streaming."""
-        
-        # Convert messages to a single prompt string
-        prompt = self._convert_messages_to_genai_format(messages)
-        
-        # Prepare generation config
-        generation_config = self._prepare_generation_config(stop)
-        
-        try:
-            # Use Google GenAI streaming
-            stream = self._client.models.generate_content_stream(
-                model=self.model_name,
-                contents=prompt,
-                config=generation_config if generation_config else None
-            )
-            
-            for chunk_response in stream:
-                if hasattr(chunk_response, 'text') and chunk_response.text:
-                    content = chunk_response.text
-                    
-                    chunk = ChatGenerationChunk(
-                        message=AIMessageChunk(
-                            content=content,
-                            response_metadata={
-                                "model_name": self.model_name,
-                                "finish_reason": getattr(chunk_response, 'finish_reason', None),
-                            }
-                        )
-                    )
-                    yield chunk
-                    
-                    # Trigger callback for new token
-                    if run_manager:
-                        run_manager.on_llm_new_token(content, chunk=chunk)
-        
-        except Exception as e:
-            # Fallback to non-streaming if streaming fails
-            try:
-                response = self._client.models.generate_content(
-                    model=self.model_name,
-                    contents=prompt,
-                    config=generation_config if generation_config else None
-                )
-                
-                generated_text = response.text if hasattr(response, 'text') else str(response)
-                
-                # Simulate streaming by yielding words
-                words = generated_text.split()
-                for i, word in enumerate(words):
-                    content = f" {word}" if i > 0 else word
-                    
-                    chunk = ChatGenerationChunk(
-                        message=AIMessageChunk(content=content)
-                    )
-                    yield chunk
-                    
-                    if run_manager:
-                        run_manager.on_llm_new_token(content, chunk=chunk)
-                        
-            except Exception as fallback_e:
-                raise ValueError(f"Error streaming content with Google GenAI: {str(e)}. Fallback also failed: {str(fallback_e)}")
-    
-    async def _astream(
-        self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> AsyncIterator[ChatGenerationChunk]:
-        """Async stream the output using Google's genai client."""
-        
-        # Convert messages to a single prompt string
-        prompt = self._convert_messages_to_genai_format(messages)
-        
-        # Prepare generation config
-        generation_config = self._prepare_generation_config(stop)
-        
-        try:
-            # Use Google GenAI streaming in async context
-            loop = asyncio.get_event_loop()
-            
-            # Run the streaming in executor
-            def create_stream():
-                return self._client.models.generate_content_stream(
-                    model=self.model_name,
-                    contents=prompt,
-                    config=generation_config if generation_config else None
-                )
-            
-            stream = await loop.run_in_executor(None, create_stream)
-            
-            for chunk_response in stream:
-                if hasattr(chunk_response, 'text') and chunk_response.text:
-                    content = chunk_response.text
-                    
-                    chunk = ChatGenerationChunk(
-                        message=AIMessageChunk(
-                            content=content,
-                            response_metadata={
-                                "model_name": self.model_name,
-                                "finish_reason": getattr(chunk_response, 'finish_reason', None),
-                            }
-                        )
-                    )
-                    yield chunk
-                    
-                    # Trigger callback for new token
-                    if run_manager:
-                        await run_manager.on_llm_new_token(content, chunk=chunk)
-        
-        except Exception as e:
-            # Fallback to async generate and simulate streaming
-            try:
-                result = await self._agenerate(messages, stop, run_manager, **kwargs)
-                generated_text = result.generations[0].message.content
-                
-                # Simulate streaming by yielding words
-                words = generated_text.split()
-                for i, word in enumerate(words):
-                    content = f" {word}" if i > 0 else word
-                    
-                    chunk = ChatGenerationChunk(
-                        message=AIMessageChunk(content=content)
-                    )
-                    yield chunk
-                    
-                    if run_manager:
-                        await run_manager.on_llm_new_token(content, chunk=chunk)
-                        
-            except Exception as fallback_e:
-                raise ValueError(f"Error async streaming content with Google GenAI: {str(e)}. Fallback also failed: {str(fallback_e)}")