rusticai-vertexai 0.0.1__tar.gz

rusticai_vertexai-0.0.1/PKG-INFO

@@ -0,0 +1,58 @@
+Metadata-Version: 2.3
+Name: rusticai-vertexai
+Version: 0.0.1
+Summary: Rustic AI module with Vertex AI integrations
+License: Apache-2.0
+Author: Dragonscale Industries Inc.
+Author-email: dev@dragonscale.ai
+Requires-Python: >=3.12,<3.13
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Provides-Extra: test
+Requires-Dist: google-cloud-aiplatform (>=1.91.0,<2.0.0)
+Requires-Dist: rusticai-core (>=0.0.8,<0.1.0)
+Project-URL: Homepage, https://www.rustic.ai/
+Project-URL: Repository, https://github.com/rustic-ai/rustic-ai
+Project-URL: Rustic AI Core, https://pypi.org/project/rusticai-core/
+Description-Content-Type: text/markdown
+
+# Rustic AI VertexAI
+
+Rustic AI module supporting integration with [Vertex AI](https://cloud.google.com/vertex-ai)
+
+## Prerequisites
+1. Using this module requires a Google Cloud Account project with Vertex AI API enabled. The account must have the role `Vertex AI User`.
+2. `gcloud CLI` for local development.
+
+## Installing
+
+```shell
+pip install rusticai-vertexai
+```
+**Note:** It depends on [rusticai-core](https://pypi.org/project/rusticai-core/)
+
+## Building from Source
+
+```shell
+poetry install --with dev
+poetry build
+```
+
+## Usage
+The library also supports using environment variables `VERTEXAI_PROJECT` and `VERTEXAI_LOCATION` in addition to the ones supported by [Google Cloud AI Platform SDK](https://cloud.google.com/python/docs/reference/aiplatform/latest) 
+```shell
+export VERTEXAI_PROJECT=<project-id>
+export VERTEXAI_LOCATION=<location>
+```
+
+For local development/use, ensure the account you are using has the roles/permissions of `Vertex AI User` and `Service Usage Consumer` at the minimum.
+Login using `gcloud CLI` in a terminal
+```shell
+gcloud auth application-default set-quota-project <project-id>
+gcloud auth application-default login
+```
+For more options/details, [refer Google Cloud docs](https://cloud.google.com/docs/authentication/set-up-adc-local-dev-environment).
+
+On Cloud, It's recommended to use Google Workload Identity.
+

rusticai_vertexai-0.0.1/README.md

@@ -0,0 +1,38 @@
+# Rustic AI VertexAI
+
+Rustic AI module supporting integration with [Vertex AI](https://cloud.google.com/vertex-ai)
+
+## Prerequisites
+1. Using this module requires a Google Cloud Account project with Vertex AI API enabled. The account must have the role `Vertex AI User`.
+2. `gcloud CLI` for local development.
+
+## Installing
+
+```shell
+pip install rusticai-vertexai
+```
+**Note:** It depends on [rusticai-core](https://pypi.org/project/rusticai-core/)
+
+## Building from Source
+
+```shell
+poetry install --with dev
+poetry build
+```
+
+## Usage
+The library also supports using environment variables `VERTEXAI_PROJECT` and `VERTEXAI_LOCATION` in addition to the ones supported by [Google Cloud AI Platform SDK](https://cloud.google.com/python/docs/reference/aiplatform/latest) 
+```shell
+export VERTEXAI_PROJECT=<project-id>
+export VERTEXAI_LOCATION=<location>
+```
+
+For local development/use, ensure the account you are using has the roles/permissions of `Vertex AI User` and `Service Usage Consumer` at the minimum.
+Login using `gcloud CLI` in a terminal
+```shell
+gcloud auth application-default set-quota-project <project-id>
+gcloud auth application-default login
+```
+For more options/details, [refer Google Cloud docs](https://cloud.google.com/docs/authentication/set-up-adc-local-dev-environment).
+
+On Cloud, It's recommended to use Google Workload Identity.

rusticai_vertexai-0.0.1/pyproject.toml

@@ -0,0 +1,72 @@
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry]
+name = "rusticai-vertexai"
+version = "0.0.1"
+description = "Rustic AI module with Vertex AI integrations"
+authors = ["Dragonscale Industries Inc. <dev@dragonscale.ai>"]
+license = "Apache-2.0"
+readme = "README.md"
+homepage = "https://www.rustic.ai/"
+repository = "https://github.com/rustic-ai/rustic-ai"
+packages = [{ include = "rustic_ai", from = "src" }]
+
+[tool.poetry.urls]
+"Rustic AI Core" = "https://pypi.org/project/rusticai-core/"
+
+[tool.poetry.dependencies]
+python = ">=3.12,<3.13"
+rusticai-core = {   version = "0.0.8"}
+google-cloud-aiplatform = "^1.91.0"
+
+[tool.poetry-monorepo.deps]
+
+
+[tool.poetry.group.dev]
+optional = true
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^8.3.4"
+black = { version = "^24.2.0", extras = ["jupyter"] }
+flake8 = "^7.1.2"
+tox = "^4.24.1"
+isort = "^6.0.0"
+mypy = "^1.15.0"
+
+[tool.poetry.extras]
+test = ["pytest"]
+
+[tool.black]
+line-length = 120
+target-version = ['py312']
+include = '\.pyi?$'
+exclude = '''
+/(
+    \.git
+  | \.mypy_cache
+  | \.tox
+  | \.venv
+  | _build
+  | buck-out
+  | build
+  | dist
+)/
+'''
+
+[tool.mypy]
+python_version = "3.12"
+ignore_missing_imports = true
+check_untyped_defs = true
+plugins = "pydantic.mypy"
+explicit_package_bases = true
+
+
+[tool.isort]
+profile = "black"
+force_sort_within_sections = true # For the sorting order issues
+lines_between_sections = 1 # For the newline issue
+known_third_party = ["pydantic"]
+known_first_party = ["rustic_ai.core", "rustic_ai.vertexai"]
+

rusticai_vertexai-0.0.1/src/rustic_ai/vertexai/agent_ext/embeddings.py

@@ -0,0 +1,88 @@
+from typing import List, Optional
+
+from pydantic.config import JsonDict
+from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
+
+from rustic_ai.core.guild.agent_ext.depends import DependencyResolver
+from rustic_ai.core.guild.agent_ext.depends.embeddings import Embeddings
+from rustic_ai.vertexai.client import VertexAIBase, VertexAIConf
+
+
+class VertexAIEmbeddingConf(VertexAIConf):
+    auto_truncate: bool = False
+    task_type: Optional[str] = None
+    output_dimensionality: Optional[int] = None
+    tokens_per_entry: int = 500
+    max_tokens_per_batch: int = 20000
+
+
+DEFAULT_MODEL: str = "text-embedding-005"
+
+
+class VertexAIEmbeddings(Embeddings, VertexAIBase):
+    def __init__(self, model: str, conf: VertexAIEmbeddingConf):
+        VertexAIBase.__init__(self, conf.project_id, conf.location)
+        self.conf = conf
+        self.model = TextEmbeddingModel.from_pretrained(model)
+
+    def _create_embedding_batches(self, text: list[str], tokens_per_entry: int) -> list[list[TextEmbeddingInput]]:
+        if not text:
+            return []
+
+        if tokens_per_entry <= 0:
+            raise ValueError("Tokens per entry must be greater than 0")
+
+        # Create batches
+        batches = []
+        current_batch = []
+        current_batch_tokens = 0
+
+        for t in text:
+            embedding_input = TextEmbeddingInput(text=t, task_type=self.conf.task_type)
+
+            # If adding this would exceed our token limit, start a new batch
+            if current_batch_tokens + tokens_per_entry < self.conf.max_tokens_per_batch:
+                # Add to current batch
+                current_batch.append(embedding_input)
+                current_batch_tokens += tokens_per_entry
+            else:
+                batches.append(current_batch)
+                current_batch = [embedding_input]
+                current_batch_tokens = tokens_per_entry
+
+        # Add the last batch if it's not empty
+        if current_batch:
+            batches.append(current_batch)
+
+        return batches
+
+    def embed(self, text: List[str]) -> List[List[float]]:
+        input_batches = self._create_embedding_batches(text, self.conf.tokens_per_entry)
+        result: List[List[float]] = []
+        for batch in input_batches:
+            embeddings = self.model.get_embeddings(
+                batch, auto_truncate=self.conf.auto_truncate, output_dimensionality=self.conf.output_dimensionality
+            )
+            for embedding in embeddings:
+                result.append(embedding.values)
+        return result
+
+
+class VertexAIEmbeddingsResolver(DependencyResolver[Embeddings]):
+    memoize_resolution: bool = False
+
+    def __init__(
+        self,
+        model_name: str = DEFAULT_MODEL,
+        conf: JsonDict = {},
+    ):
+        super().__init__()
+        embedding_conf = VertexAIEmbeddingConf.model_validate(conf)
+
+        self.embedding = VertexAIEmbeddings(
+            model=model_name,
+            conf=embedding_conf,
+        )
+
+    def resolve(self, guild_id: str, agent_id: str) -> Embeddings:
+        return self.embedding

rusticai_vertexai-0.0.1/src/rustic_ai/vertexai/agents/image_generation.py

@@ -0,0 +1,81 @@
+import mimetypes
+from typing import Literal, Optional
+import uuid
+
+from pydantic import BaseModel
+from vertexai.vision_models import ImageGenerationModel
+
+from rustic_ai.core import Agent, AgentMode, AgentSpec, AgentType
+from rustic_ai.core.agents.commons.image_generation import ImageGenerationResponse
+from rustic_ai.core.agents.commons.media import MediaLink
+from rustic_ai.core.guild import agent
+from rustic_ai.core.guild.agent_ext.depends.filesystem import FileSystem
+from rustic_ai.core.guild.dsl import BaseAgentProps
+from rustic_ai.vertexai.client import VertexAIBase, VertexAIConf
+
+
+class VertexAiImagenAgentProps(BaseAgentProps, VertexAIConf):
+    model_id: str = "imagen-3.0-fast-generate-001"
+    add_watermark: Optional[bool] = True
+    safety_filter_level: Optional[Literal["block_most", "block_some", "block_few", "block_fewest"]] = None
+    person_generation: Optional[Literal["dont_allow", "allow_adult", "allow_all"]] = None
+
+
+class VertexAiImageGenerationRequest(BaseModel):
+    prompt: str
+    negative_prompt: Optional[str] = None
+    number_of_images: int = 1
+    aspect_ratio: Optional[Literal["1:1", "9:16", "16:9", "4:3", "3:4"]] = None
+    guidance_scale: Optional[float] = None
+    language: Optional[str] = None
+    seed: Optional[int] = None
+    image_format: Optional[str] = "png"
+
+
+class VertexAiImagenAgent(Agent[VertexAiImagenAgentProps], VertexAIBase):
+    def __init__(self, agent_spec: AgentSpec[VertexAiImagenAgentProps]):
+        if agent_spec.properties is None:
+            agent_spec.properties = VertexAiImagenAgentProps()
+        Agent.__init__(
+            self,
+            agent_spec=agent_spec,
+            agent_type=AgentType.BOT,
+            agent_mode=AgentMode.REMOTE,
+        )
+
+        self.conf = agent_spec.properties
+        VertexAIBase.__init__(self, self.conf.project_id, self.conf.location)
+        print("initialized vertexai")
+        self.model = ImageGenerationModel.from_pretrained(self.conf.model_id)
+
+    @agent.processor(VertexAiImageGenerationRequest, depends_on=["filesystem:guild_fs:True"])
+    def generate_image(self, ctx: agent.ProcessContext[VertexAiImageGenerationRequest], guild_fs: FileSystem) -> None:
+        image_gen_request = ctx.payload
+
+        output_images = self.model.generate_images(
+            prompt=image_gen_request.prompt,
+            negative_prompt=image_gen_request.negative_prompt,
+            number_of_images=image_gen_request.number_of_images,
+            aspect_ratio=image_gen_request.aspect_ratio,
+            guidance_scale=image_gen_request.guidance_scale,
+            language=image_gen_request.language,
+            seed=image_gen_request.seed,
+            add_watermark=self.conf.add_watermark,
+            safety_filter_level=self.conf.safety_filter_level,
+            person_generation=self.conf.person_generation,
+        ).images
+
+        result = ImageGenerationResponse(files=[], errors=[], request=image_gen_request.model_dump_json())
+        for i, image in enumerate(output_images):
+            filename = f"{uuid.uuid4()}.{image_gen_request.image_format}"
+            try:
+                with guild_fs.open(filename, "wb") as f:
+                    image.save(f)
+                # Create a MediaLink object for the image
+                media_link = MediaLink(
+                    url=filename, name=filename, mimetype=mimetypes.guess_type(filename)[0], on_filesystem=True
+                )
+                result.files.append(media_link)
+            except Exception as e:
+                result.errors.append(f"Failed to write image file {filename}:{e}")
+        ctx.send(result)

rusticai_vertexai-0.0.1/src/rustic_ai/vertexai/client.py

@@ -0,0 +1,63 @@
+import os
+from typing import Optional
+
+from google.cloud import aiplatform
+from pydantic import BaseModel, ConfigDict
+
+
+class VertexAIConf(BaseModel):
+    """Configuration class for Vertex AI.
+
+    This class is used to encapsulate configuration details required for
+    interacting with Vertex AI.
+
+    Attributes:
+        project_id (Optional[str]): The Google Cloud project ID associated with
+            Vertex AI.
+        location (Optional[str]): The region or location where the Vertex AI
+            resources are hosted.
+    """
+
+    project_id: Optional[str] = None
+    location: Optional[str] = None
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+
+class VertexAIBase:
+    """Base class for Google Cloud Vertex AI services.
+
+    This class handles the common initialization pattern for the Vertex AI SDK,
+    ensuring that it's only initialized once per runtime. Other Vertex AI
+    service implementations can inherit from this class to reuse this logic.
+    """
+
+    _is_vertexai_initialized = False
+
+    def __init__(
+        self,
+        project_id: Optional[str] = None,
+        location: Optional[str] = None,
+    ):
+        """Initialize Vertex AI SDK if not already initialized.
+
+        Args:
+            project_id (Optional[str]): The Google Cloud project ID. If None,
+                attempts to use the VERTEXAI_PROJECT environment variable.
+            location (Optional[str]): The region or location where the Vertex AI
+                resources are hosted. If None, attempts to use the VERTEXAI_LOCATION
+                environment variable.
+
+        Raises:
+            ValueError: If project_id or location cannot be determined from
+                parameters or environment variables.
+        """
+        if not self._is_vertexai_initialized:
+            if project_id is None:
+                project_id = os.environ.get("VERTEXAI_PROJECT")
+
+            if location is None:
+                location = os.environ.get("VERTEXAI_LOCATION")
+
+            aiplatform.init(project=project_id, location=location)
+            self._is_vertexai_initialized = True