local-coze 0.0.1__py3-none-any.whl

local_coze/embedding/client.py

@@ -0,0 +1,350 @@
+import asyncio
+from typing import Optional, List, Dict
+from cozeloop.decorator import observe
+from coze_coding_utils.runtime_ctx.context import Context
+
+from ..core.client import BaseClient
+from ..core.config import Config
+from ..core.exceptions import APIError, ValidationError
+from .models import (
+    EmbeddingConfig,
+    EmbeddingInputItem,
+    EmbeddingInputImageURL,
+    EmbeddingInputVideoURL,
+    EmbeddingRequest,
+    EmbeddingResponse,
+    EmbeddingData,
+    EmbeddingUsage,
+    MultiEmbeddingConfig,
+    SparseEmbeddingConfig,
+    SparseEmbeddingItem,
+)
+
+
+class EmbeddingClient(BaseClient):
+    def __init__(
+        self,
+        config: Optional[Config] = None,
+        ctx: Optional[Context] = None,
+        custom_headers: Optional[Dict[str, str]] = None,
+        verbose: bool = False,
+    ):
+        super().__init__(config, ctx, custom_headers, verbose)
+        self.base_url = self.config.base_url
+        self.model = EmbeddingConfig.DEFAULT_MODEL
+
+    def _validate_inputs(
+        self,
+        texts: Optional[List[str]],
+        image_urls: Optional[List[str]],
+        video_urls: Optional[List[str]] = None
+    ) -> None:
+        has_texts = texts is not None and len(texts) > 0
+        has_images = image_urls is not None and len(image_urls) > 0
+        has_videos = video_urls is not None and len(video_urls) > 0
+
+        if not has_texts and not has_images and not has_videos:
+            raise ValidationError(
+                "At least one of texts, image_urls, or video_urls must be provided",
+                field="input"
+            )
+
+        total_inputs = 0
+        if texts:
+            total_inputs += len(texts)
+        if image_urls:
+            total_inputs += len(image_urls)
+        if video_urls:
+            total_inputs += len(video_urls)
+
+        if total_inputs > EmbeddingConfig.MAX_BATCH_SIZE:
+            raise ValidationError(
+                f"Total inputs exceed maximum batch size of {EmbeddingConfig.MAX_BATCH_SIZE}",
+                field="input",
+                value=total_inputs
+            )
+
+    def _build_input_items(
+        self,
+        texts: Optional[List[str]],
+        image_urls: Optional[List[str]],
+        video_urls: Optional[List[str]] = None
+    ) -> List[EmbeddingInputItem]:
+        items = []
+        if texts:
+            for text in texts:
+                if text and text.strip():
+                    items.append(EmbeddingInputItem(type="text", text=text))
+        if image_urls:
+            for url in image_urls:
+                if url and url.strip():
+                    items.append(EmbeddingInputItem(
+                        type="image_url",
+                        image_url=EmbeddingInputImageURL(url=url)
+                    ))
+        if video_urls:
+            for url in video_urls:
+                if url and url.strip():
+                    items.append(EmbeddingInputItem(
+                        type="video_url",
+                        video_url=EmbeddingInputVideoURL(url=url)
+                    ))
+        return items
+
+    @observe
+    def embed(
+        self,
+        texts: Optional[List[str]] = None,
+        image_urls: Optional[List[str]] = None,
+        video_urls: Optional[List[str]] = None,
+        model: Optional[str] = None,
+        dimensions: Optional[int] = None,
+        encoding_format: Optional[str] = None,
+        instructions: Optional[str] = None,
+        multi_embedding: bool = False,
+        sparse_embedding: bool = False
+    ) -> EmbeddingResponse:
+        self._validate_inputs(texts, image_urls, video_urls)
+
+        input_items = self._build_input_items(texts, image_urls, video_urls)
+        if not input_items:
+            raise ValidationError("No valid input provided", field="input")
+
+        request = EmbeddingRequest(
+            model=model or self.model,
+            input=input_items,
+            dimensions=dimensions,
+            encoding_format=encoding_format,
+            instructions=instructions,
+            multi_embedding=MultiEmbeddingConfig(type="enabled") if multi_embedding else None,
+            sparse_embedding=SparseEmbeddingConfig(type="enabled") if sparse_embedding else None
+        )
+
+        data = self._request(
+            method="POST",
+            url=f"{self.base_url}/api/v3/embeddings/multimodal",
+            json=request.to_api_request()
+        )
+
+        if "error" in data and data["error"]:
+            raise APIError(
+                f"API returned error: {data['error'].get('message', 'Unknown error')}",
+                code=data['error'].get('code')
+            )
+
+        return EmbeddingResponse(**data)
+
+    @observe
+    def embed_text(
+        self,
+        text: str,
+        model: Optional[str] = None,
+        dimensions: Optional[int] = None,
+        instructions: Optional[str] = None
+    ) -> List[float]:
+        response = self.embed(
+            texts=[text],
+            model=model,
+            dimensions=dimensions,
+            instructions=instructions
+        )
+        if response.embedding:
+            return response.embedding
+        raise APIError("No embedding returned")
+
+    @observe
+    def embed_texts(
+        self,
+        texts: List[str],
+        model: Optional[str] = None,
+        dimensions: Optional[int] = None,
+        instructions: Optional[str] = None
+    ) -> List[float]:
+        response = self.embed(
+            texts=texts,
+            model=model,
+            dimensions=dimensions,
+            instructions=instructions
+        )
+        if response.embedding:
+            return response.embedding
+        raise APIError("No embedding returned")
+
+    @observe
+    def embed_image(
+        self,
+        image_url: str,
+        model: Optional[str] = None,
+        dimensions: Optional[int] = None
+    ) -> List[float]:
+        response = self.embed(
+            image_urls=[image_url],
+            model=model,
+            dimensions=dimensions
+        )
+        if response.embedding:
+            return response.embedding
+        raise APIError("No embedding returned")
+
+    @observe
+    def embed_images(
+        self,
+        image_urls: List[str],
+        model: Optional[str] = None,
+        dimensions: Optional[int] = None
+    ) -> List[float]:
+        response = self.embed(
+            image_urls=image_urls,
+            model=model,
+            dimensions=dimensions
+        )
+        if response.embedding:
+            return response.embedding
+        raise APIError("No embedding returned")
+
+    @observe
+    def embed_video(
+        self,
+        video_url: str,
+        model: Optional[str] = None,
+        dimensions: Optional[int] = None
+    ) -> List[float]:
+        response = self.embed(
+            video_urls=[video_url],
+            model=model,
+            dimensions=dimensions
+        )
+        if response.embedding:
+            return response.embedding
+        raise APIError("No embedding returned")
+
+    @observe
+    def embed_videos(
+        self,
+        video_urls: List[str],
+        model: Optional[str] = None,
+        dimensions: Optional[int] = None
+    ) -> List[float]:
+        response = self.embed(
+            video_urls=video_urls,
+            model=model,
+            dimensions=dimensions
+        )
+        if response.embedding:
+            return response.embedding
+        raise APIError("No embedding returned")
+
+    @observe
+    def embed_multimodal(
+        self,
+        texts: Optional[List[str]] = None,
+        image_urls: Optional[List[str]] = None,
+        video_urls: Optional[List[str]] = None,
+        model: Optional[str] = None,
+        dimensions: Optional[int] = None,
+        instructions: Optional[str] = None,
+        multi_embedding: bool = False,
+        sparse_embedding: bool = False
+    ) -> EmbeddingResponse:
+        return self.embed(
+            texts=texts,
+            image_urls=image_urls,
+            video_urls=video_urls,
+            model=model,
+            dimensions=dimensions,
+            instructions=instructions,
+            multi_embedding=multi_embedding,
+            sparse_embedding=sparse_embedding
+        )
+
+    @observe
+    def embed_with_multi_vectors(
+        self,
+        texts: Optional[List[str]] = None,
+        image_urls: Optional[List[str]] = None,
+        video_urls: Optional[List[str]] = None,
+        model: Optional[str] = None,
+        dimensions: Optional[int] = None,
+        instructions: Optional[str] = None
+    ) -> Optional[List[List[float]]]:
+        response = self.embed(
+            texts=texts,
+            image_urls=image_urls,
+            video_urls=video_urls,
+            model=model,
+            dimensions=dimensions,
+            instructions=instructions,
+            multi_embedding=True
+        )
+        return response.multi_embeddings
+
+    @observe
+    def embed_with_sparse(
+        self,
+        texts: Optional[List[str]] = None,
+        image_urls: Optional[List[str]] = None,
+        video_urls: Optional[List[str]] = None,
+        model: Optional[str] = None,
+        dimensions: Optional[int] = None,
+        instructions: Optional[str] = None
+    ) -> Optional[List[SparseEmbeddingItem]]:
+        response = self.embed(
+            texts=texts,
+            image_urls=image_urls,
+            video_urls=video_urls,
+            model=model,
+            dimensions=dimensions,
+            instructions=instructions,
+            sparse_embedding=True
+        )
+        return response.sparse_embeddings
+
+    async def embed_async(
+        self,
+        texts: Optional[List[str]] = None,
+        image_urls: Optional[List[str]] = None,
+        video_urls: Optional[List[str]] = None,
+        model: Optional[str] = None,
+        dimensions: Optional[int] = None,
+        encoding_format: Optional[str] = None,
+        instructions: Optional[str] = None,
+        multi_embedding: bool = False,
+        sparse_embedding: bool = False
+    ) -> EmbeddingResponse:
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(
+            None,
+            lambda: self.embed(
+                texts=texts,
+                image_urls=image_urls,
+                video_urls=video_urls,
+                model=model,
+                dimensions=dimensions,
+                encoding_format=encoding_format,
+                instructions=instructions,
+                multi_embedding=multi_embedding,
+                sparse_embedding=sparse_embedding
+            )
+        )
+
+    async def batch_embed(
+        self,
+        text_batches: List[List[str]],
+        model: Optional[str] = None,
+        dimensions: Optional[int] = None,
+        instructions: Optional[str] = None,
+        max_concurrent: int = 5
+    ) -> List[EmbeddingResponse]:
+        semaphore = asyncio.Semaphore(max_concurrent)
+
+        async def _embed_with_semaphore(texts: List[str]) -> EmbeddingResponse:
+            async with semaphore:
+                return await self.embed_async(
+                    texts=texts,
+                    model=model,
+                    dimensions=dimensions,
+                    instructions=instructions
+                )
+
+        tasks = [_embed_with_semaphore(batch) for batch in text_batches]
+        return await asyncio.gather(*tasks, return_exceptions=False)

local_coze/embedding/models.py

@@ -0,0 +1,130 @@
+from typing import Optional, List, Literal
+from pydantic import BaseModel, Field
+
+
+class EmbeddingConfig:
+    DEFAULT_MODEL = "doubao-embedding-vision-251215"
+    DEFAULT_ENCODING_FORMAT = "float"
+    MAX_BATCH_SIZE = 100
+
+
+class EmbeddingInputImageURL(BaseModel):
+    url: str = Field(..., description="Image URL")
+
+
+class EmbeddingInputVideoURL(BaseModel):
+    url: str = Field(..., description="Video URL")
+
+
+class EmbeddingInputItem(BaseModel):
+    type: Literal["text", "image_url", "video_url"] = Field(..., description="Input type")
+    text: Optional[str] = Field(default=None, description="Text content for embedding")
+    image_url: Optional[EmbeddingInputImageURL] = Field(default=None, description="Image URL for embedding")
+    video_url: Optional[EmbeddingInputVideoURL] = Field(default=None, description="Video URL for embedding")
+
+
+class MultiEmbeddingConfig(BaseModel):
+    type: Literal["enabled", "disabled"] = Field(default="disabled", description="Multi embedding mode")
+
+
+class SparseEmbeddingConfig(BaseModel):
+    type: Literal["enabled", "disabled"] = Field(default="disabled", description="Sparse embedding mode")
+
+
+class SparseEmbeddingItem(BaseModel):
+    index: int = Field(..., description="Token index")
+    value: float = Field(..., description="Token weight value")
+
+
+class PromptTokensDetails(BaseModel):
+    image_tokens: int = Field(default=0, description="Number of image tokens")
+    text_tokens: int = Field(default=0, description="Number of text tokens")
+
+
+class EmbeddingUsage(BaseModel):
+    prompt_tokens: int = Field(default=0, description="Number of prompt tokens")
+    prompt_tokens_details: Optional[PromptTokensDetails] = Field(default=None, description="Token details breakdown")
+    total_tokens: int = Field(default=0, description="Total number of tokens")
+
+
+class EmbeddingData(BaseModel):
+    object: str = Field(default="embedding", description="Object type")
+    embedding: Optional[List[float]] = Field(default=None, description="Embedding vector")
+    multi_embedding: Optional[List[List[float]]] = Field(default=None, description="Multi embedding vectors")
+    sparse_embedding: Optional[List[SparseEmbeddingItem]] = Field(default=None, description="Sparse embedding")
+    index: int = Field(default=0, description="Index of the embedding in the batch")
+
+
+class EmbeddingRequest(BaseModel):
+    model: str = Field(default=EmbeddingConfig.DEFAULT_MODEL, description="Model ID for embedding")
+    input: List[EmbeddingInputItem] = Field(..., description="List of inputs to embed")
+    dimensions: Optional[int] = Field(default=None, description="Output embedding dimensions")
+    encoding_format: Optional[Literal["float", "base64"]] = Field(default=None, description="Encoding format")
+    instructions: Optional[str] = Field(default=None, description="Instructions for embedding")
+    multi_embedding: Optional[MultiEmbeddingConfig] = Field(default=None, description="Multi embedding config")
+    sparse_embedding: Optional[SparseEmbeddingConfig] = Field(default=None, description="Sparse embedding config")
+
+    def to_api_request(self) -> dict:
+        request_data = {
+            "model": self.model,
+            "input": []
+        }
+        for item in self.input:
+            item_dict = {"type": item.type}
+            if item.text is not None:
+                item_dict["text"] = item.text
+            if item.image_url is not None:
+                item_dict["image_url"] = {"url": item.image_url.url}
+            if item.video_url is not None:
+                item_dict["video_url"] = {"url": item.video_url.url}
+            request_data["input"].append(item_dict)
+
+        if self.dimensions is not None:
+            request_data["dimensions"] = self.dimensions
+        if self.encoding_format is not None:
+            request_data["encoding_format"] = self.encoding_format
+        if self.instructions is not None:
+            request_data["instructions"] = self.instructions
+        if self.multi_embedding is not None:
+            request_data["multi_embedding"] = {"type": self.multi_embedding.type}
+        if self.sparse_embedding is not None:
+            request_data["sparse_embedding"] = {"type": self.sparse_embedding.type}
+        return request_data
+
+
+class EmbeddingResponse(BaseModel):
+    object: str = Field(default="list", description="Response object type")
+    data: Optional[EmbeddingData] = Field(default=None, description="Embedding result")
+    model: str = Field(default="", description="Model used for embedding")
+    usage: Optional[EmbeddingUsage] = Field(default=None, description="Token usage information")
+    id: Optional[str] = Field(default=None, description="Response ID")
+    created: Optional[int] = Field(default=None, description="Creation timestamp")
+    error: Optional[dict] = Field(default=None, description="Error information if any")
+
+    @property
+    def success(self) -> bool:
+        return self.error is None and self.data is not None
+
+    @property
+    def embedding(self) -> Optional[List[float]]:
+        if self.data and self.data.embedding:
+            return self.data.embedding
+        return None
+
+    @property
+    def multi_embeddings(self) -> Optional[List[List[float]]]:
+        if self.data and self.data.multi_embedding:
+            return self.data.multi_embedding
+        return None
+
+    @property
+    def sparse_embeddings(self) -> Optional[List[SparseEmbeddingItem]]:
+        if self.data and self.data.sparse_embedding:
+            return self.data.sparse_embedding
+        return None
+
+    @property
+    def error_message(self) -> Optional[str]:
+        if self.error:
+            return self.error.get("message", "Unknown error")
+        return None

local_coze/image/__init__.py

@@ -0,0 +1,19 @@
+from .client import ImageGenerationClient
+from .models import (
+    ImageConfig,
+    ImageSize,
+    ImageGenerationRequest,
+    ImageGenerationResponse,
+    ImageData,
+    UsageInfo
+)
+
+__all__ = [
+    "ImageGenerationClient",
+    "ImageConfig",
+    "ImageSize",
+    "ImageGenerationRequest",
+    "ImageGenerationResponse",
+    "ImageData",
+    "UsageInfo",
+]

local_coze/image/client.py

@@ -0,0 +1,110 @@
+import asyncio
+from typing import Any, Dict, List, Optional, Union
+
+from coze_coding_utils.runtime_ctx.context import Context
+from cozeloop.decorator import observe
+
+from ..core.client import BaseClient
+from ..core.config import Config
+from ..core.exceptions import APIError
+from .models import ImageConfig, ImageGenerationRequest, ImageGenerationResponse
+
+
+class ImageGenerationClient(BaseClient):
+    def __init__(
+        self,
+        config: Optional[Config] = None,
+        ctx: Optional[Context] = None,
+        custom_headers: Optional[Dict[str, str]] = None,
+        verbose: bool = False,
+    ):
+        super().__init__(config, ctx, custom_headers, verbose)
+        self.base_url = self.config.base_url
+        self.model = ImageConfig.DEFAULT_MODEL
+
+    def extract_urls(self, response: ImageGenerationResponse) -> List[str]:
+        urls = []
+        for item in response.data:
+            if item.error:
+                raise APIError(
+                    f"图片生成失败: {item.error.get('message', 'Unknown error')}",
+                    code=item.error.get("code"),
+                )
+            if item.url:
+                urls.append(item.url)
+            elif item.b64_json:
+                urls.append(f"data:image/png;base64,{item.b64_json}")
+        return urls
+
+    @observe
+    def generate(
+        self,
+        prompt: str,
+        size: Optional[str] = None,
+        watermark: Optional[bool] = None,
+        image: Optional[Union[str, List[str]]] = None,
+        response_format: Optional[str] = None,
+        optimize_prompt_mode: Optional[str] = None,
+        sequential_image_generation: Optional[str] = None,
+        sequential_image_generation_max_images: Optional[int] = None,
+    ) -> ImageGenerationResponse:
+        request_params = {"prompt": prompt}
+        if size is not None:
+            request_params["size"] = size
+        if watermark is not None:
+            request_params["watermark"] = watermark
+        if image is not None:
+            request_params["image"] = image
+        if response_format is not None:
+            request_params["response_format"] = response_format
+        if optimize_prompt_mode is not None:
+            request_params["optimize_prompt_mode"] = optimize_prompt_mode
+        if sequential_image_generation is not None:
+            request_params["sequential_image_generation"] = sequential_image_generation
+        if sequential_image_generation_max_images is not None:
+            request_params["sequential_image_generation_max_images"] = (
+                sequential_image_generation_max_images
+            )
+
+        request = ImageGenerationRequest(**request_params)
+
+        data = self._request(
+            method="POST",
+            url=f"{self.base_url}/api/v3/images/generations",
+            json=request.to_api_request(self.model),
+        )
+
+        if "error" in data and data["error"]:
+            raise APIError(
+                f"API 返回错误: {data['error'].get('message', 'Unknown error')}",
+                code=data["error"].get("code"),
+            )
+
+        parsed_response = ImageGenerationResponse(**data)
+
+        return parsed_response
+
+    async def generate_async(
+        self,
+        prompt: str,
+        size: Optional[str] = None,
+        watermark: Optional[bool] = None,
+        image: Optional[Union[str, List[str]]] = None,
+        response_format: Optional[str] = None,
+        optimize_prompt_mode: Optional[str] = None,
+        sequential_image_generation: Optional[str] = None,
+        sequential_image_generation_max_images: Optional[int] = None,
+    ) -> ImageGenerationResponse:
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(
+            None,
+            self.generate,
+            prompt,
+            size,
+            watermark,
+            image,
+            response_format,
+            optimize_prompt_mode,
+            sequential_image_generation,
+            sequential_image_generation_max_images,
+        )