local-coze 0.0.1__py3-none-any.whl

local_coze/image/models.py

@@ -0,0 +1,163 @@
+import re
+from typing import Optional, Union, List, Literal
+from pydantic import BaseModel, Field, field_validator
+
+from ..core.exceptions import ValidationError
+
+
+class ImageConfig:
+    DEFAULT_MODEL = "doubao-seedream-4-5-251128"
+    DEFAULT_SIZE = "2K"
+    DEFAULT_CUSTOM_SIZE = "2048x2048"
+    DEFAULT_WATERMARK = True
+    DEFAULT_RESPONSE_FORMAT = "url"
+    DEFAULT_OPTIMIZE_PROMPT_MODE = "standard"
+    DEFAULT_SEQUENTIAL_IMAGE_GENERATION = "disabled"
+    DEFAULT_SEQUENTIAL_IMAGE_GENERATION_MAX_IMAGES = 15
+    
+    MIN_TOTAL_PIXELS = 2560 * 1440
+    MAX_TOTAL_PIXELS = 4096 * 4096
+    MIN_ASPECT_RATIO = 1 / 16
+    MAX_ASPECT_RATIO = 16
+    
+    PRESET_SIZES = ["2K", "4K"]
+
+
+class ImageSize:
+    @staticmethod
+    def validate(size: str) -> str:
+        if size in ImageConfig.PRESET_SIZES:
+            return size
+        
+        match = re.match(r'^(\d+)x(\d+)$', size)
+        if not match:
+            raise ValidationError(
+                f"尺寸格式必须是 {ImageConfig.PRESET_SIZES} 或 'WIDTHxHEIGHT' 格式（如 2048x2048）",
+                field="size",
+                value=size
+            )
+        
+        width, height = int(match.group(1)), int(match.group(2))
+        
+        if width <= 0 or height <= 0:
+            raise ValidationError(
+                "宽度和高度必须为正整数",
+                field="size",
+                value=size
+            )
+        
+        total_pixels = width * height
+        if not (ImageConfig.MIN_TOTAL_PIXELS <= total_pixels <= ImageConfig.MAX_TOTAL_PIXELS):
+            raise ValidationError(
+                f"总像素数必须在 [{ImageConfig.MIN_TOTAL_PIXELS:,}, {ImageConfig.MAX_TOTAL_PIXELS:,}] 范围内，"
+                f"当前值: {total_pixels:,} ({width}x{height})",
+                field="size",
+                value=size
+            )
+        
+        aspect_ratio = width / height
+        if not (ImageConfig.MIN_ASPECT_RATIO <= aspect_ratio <= ImageConfig.MAX_ASPECT_RATIO):
+            raise ValidationError(
+                f"宽高比必须在 [1:16, 16:1] 范围内，当前值: {width}:{height} ({aspect_ratio:.2f})",
+                field="size",
+                value=size
+            )
+        
+        return size
+    
+    @staticmethod
+    def validate_or_default(size: str) -> str:
+        try:
+            return ImageSize.validate(size)
+        except ValidationError:
+            return ImageConfig.DEFAULT_SIZE
+
+
+class ImageGenerationRequest(BaseModel):
+    prompt: str = Field(..., description="用于生成图像的提示词")
+    size: str = Field(default=ImageConfig.DEFAULT_SIZE, description="生成图片的尺寸")
+    watermark: bool = Field(default=ImageConfig.DEFAULT_WATERMARK, description="是否添加水印")
+    image: Optional[Union[str, List[str]]] = Field(default=None, description="参考图片")
+    response_format: Literal["url", "b64_json"] = Field(
+        default=ImageConfig.DEFAULT_RESPONSE_FORMAT,
+        description="返回格式"
+    )
+    optimize_prompt_mode: str = Field(
+        default=ImageConfig.DEFAULT_OPTIMIZE_PROMPT_MODE,
+        description="提示词优化模式"
+    )
+    sequential_image_generation: Literal["auto", "disabled"] = Field(
+        default=ImageConfig.DEFAULT_SEQUENTIAL_IMAGE_GENERATION,
+        description="组图功能"
+    )
+    sequential_image_generation_max_images: int = Field(
+        default=ImageConfig.DEFAULT_SEQUENTIAL_IMAGE_GENERATION_MAX_IMAGES,
+        ge=1,
+        le=15,
+        description="最大图片数量"
+    )
+    
+    @field_validator('size')
+    @classmethod
+    def validate_size(cls, v: str) -> str:
+        return ImageSize.validate_or_default(v)
+    
+    def to_api_request(self, model: str) -> dict:
+        return {
+            "model": model,
+            "prompt": self.prompt,
+            "size": self.size,
+            "watermark": self.watermark,
+            "image": self.image,
+            "response_format": self.response_format,
+            "optimize_prompt_options": {
+                "mode": self.optimize_prompt_mode,
+            },
+            "sequential_image_generation": self.sequential_image_generation,
+            "sequential_image_generation_options": {
+                "max_images": self.sequential_image_generation_max_images,
+            },
+        }
+
+
+class ImageData(BaseModel):
+    url: Optional[str] = None
+    b64_json: Optional[str] = None
+    size: Optional[str] = None
+    error: Optional[dict] = None
+
+
+class UsageInfo(BaseModel):
+    generated_images: int = 0
+    output_tokens: Optional[int] = None
+    total_tokens: Optional[int] = None
+
+
+class ImageGenerationResponse(BaseModel):
+    model: str
+    created: int
+    data: List[ImageData]
+    usage: Optional[UsageInfo] = None
+    error: Optional[dict] = None
+    
+    @property
+    def success(self) -> bool:
+        return self.error is None and all(item.error is None for item in self.data)
+    
+    @property
+    def image_urls(self) -> List[str]:
+        return [item.url for item in self.data if item.url]
+    
+    @property
+    def image_b64_list(self) -> List[str]:
+        return [item.b64_json for item in self.data if item.b64_json]
+    
+    @property
+    def error_messages(self) -> List[str]:
+        messages = []
+        if self.error:
+            messages.append(f"API错误: {self.error.get('message', 'Unknown error')}")
+        for item in self.data:
+            if item.error:
+                messages.append(f"图片生成错误: {item.error.get('message', 'Unknown error')}")
+        return messages

local_coze/knowledge/__init__.py

@@ -0,0 +1,19 @@
+from .client import KnowledgeClient
+from .models import (
+    ChunkConfig,
+    DataSourceType,
+    KnowledgeChunk,
+    KnowledgeDocument,
+    KnowledgeInsertResponse,
+    KnowledgeSearchResponse,
+)
+
+__all__ = [
+    "KnowledgeClient",
+    "ChunkConfig",
+    "DataSourceType",
+    "KnowledgeSearchResponse",
+    "KnowledgeInsertResponse",
+    "KnowledgeChunk",
+    "KnowledgeDocument",
+]

local_coze/knowledge/client.py

@@ -0,0 +1,148 @@
+from typing import Dict, List, Optional, Union
+
+from coze_coding_utils.runtime_ctx.context import Context
+
+from ..core.client import BaseClient
+from ..core.config import Config
+from .models import (
+    ChunkConfig,
+    KnowledgeChunk,
+    KnowledgeDocument,
+    KnowledgeInsertResponse,
+    KnowledgeSearchResponse,
+)
+
+
+class KnowledgeClient(BaseClient):
+    def __init__(
+        self,
+        config: Optional[Config] = None,
+        ctx: Optional[Context] = None,
+        custom_headers: Optional[Dict[str, str]] = None,
+        verbose: bool = True,
+    ):
+        super().__init__(config, ctx, custom_headers, verbose)
+        # Use COZE_INTEGRATION_BASE_URL as endpoint
+        self.base_url = self.config.base_url
+
+    def search(
+        self,
+        query: str,
+        table_names: Optional[List[str]] = None,
+        top_k: int = 5,
+        min_score: Optional[float] = 0.0,
+        extra_headers: Optional[Dict[str, str]] = None,
+    ) -> KnowledgeSearchResponse:
+        """
+        Search for knowledge chunks in specified tables.
+
+        Args:
+            query: The search query string.
+            table_names: List of table names to search in (dataset).
+            top_k: Number of results to return. Default is 5.
+            min_score: Minimum similarity score. Default is 0.0.
+            extra_headers: Extra headers to send with the request.
+
+        Returns:
+            KnowledgeSearchResponse: The search results.
+        """
+        payload = {
+            "query": query,
+            "top_k": top_k,
+            "min_score": min_score,
+        }
+        if table_names:
+            payload["dataset"] = table_names
+        
+        url = f"{self.base_url}/v1/knowledge_base/recall"
+
+        headers = extra_headers or {}
+        if self.config.api_key:
+             headers["x-coze-token"] = f"Bearer {self.config.api_key}"
+
+        response = self._request(
+            method="POST",
+            url=url,
+            json=payload,
+            headers=headers,
+        )
+
+        # Response structure: { "data": [ ... ], "BaseResp": ... }
+        # RecallDataInfo: { "slice": "...", "score": 0.0, "chunk_id": "...", "doc_id": "..." }
+        
+        data_list = response.get("data", [])
+        
+        chunks = []
+        for item in data_list:
+            # Map RecallDataInfo to KnowledgeChunk
+            chunks.append(KnowledgeChunk(
+                content=item.get("slice", ""),
+                score=item.get("score", 0.0),
+                chunk_id=item.get("chunk_id"),
+                doc_id=item.get("doc_id")
+            ))
+            
+        return KnowledgeSearchResponse(
+            chunks=chunks,
+            code=response.get("code", 0),
+            msg=response.get("msg", "")
+        )
+
+    def add_documents(
+        self,
+        documents: List[Union[KnowledgeDocument, Dict]],
+        table_name: str,
+        chunk_config: Optional[ChunkConfig] = None,
+        extra_headers: Optional[Dict[str, str]] = None,
+    ) -> KnowledgeInsertResponse:
+        """
+        Add documents to a specified table in the knowledge base (BatchImportData).
+
+        Args:
+            documents: List of documents to add. Can be KnowledgeDocument objects or dictionaries.
+            table_name: The name of the table to add documents to (dataset).
+            chunk_config: Optional chunking configuration.
+            extra_headers: Extra headers to send with the request.
+
+        Returns:
+            KnowledgeInsertResponse: The response containing inserted document IDs and status.
+        """
+        docs_payload = []
+        for doc in documents:
+            if isinstance(doc, KnowledgeDocument):
+                docs_payload.append(doc.to_api_format())
+            else:
+                # Assuming dictionary is already in correct format or close to it
+                docs_payload.append(doc)
+
+        payload = {
+            "dataset": table_name,
+            "data": docs_payload,
+        }
+        
+        if chunk_config:
+            payload["chunk_config"] = chunk_config.model_dump(exclude_none=True)
+
+        url = f"{self.base_url}/v1/knowledge_base/batch_import"
+
+        headers = extra_headers or {}
+        if self.config.api_key:
+             headers["x-coze-token"] = f"Bearer {self.config.api_key}"
+
+        response = self._request(
+            method="POST",
+            url=url,
+            json=payload,
+            headers=headers,
+        )
+        
+        code = response.get("code", 0)
+        msg = response.get("msg", "")
+        if code == 0:
+            msg = "成功！文档正在异步导入中，请稍候。您可以查看数据库中的 'knowledge' schema 来验证导入状态。"
+
+        return KnowledgeInsertResponse(
+            doc_ids=response.get("doc_ids"),
+            code=code,
+            msg=msg
+        )

local_coze/knowledge/models.py

@@ -0,0 +1,45 @@
+from enum import IntEnum
+from typing import List, Optional, Dict, Any, Union
+from pydantic import BaseModel, Field
+
+
+class DataSourceType(IntEnum):
+    TEXT = 0
+    URL = 1
+    URI = 2
+
+
+class ChunkConfig(BaseModel):
+    separator: str = Field(..., description="Chunk separator")
+    max_tokens: int = Field(..., description="Max tokens per chunk")
+    remove_extra_spaces: bool = Field(False, description="Normalize extra spaces")
+    remove_urls_emails: bool = Field(False, description="Strip URLs and emails")
+
+
+class KnowledgeDocument(BaseModel):
+    source: DataSourceType = Field(..., description="Data source type")
+    raw_data: Optional[str] = Field(None, description="Plain text content (for RAW_TEXT)")
+    uri: Optional[str] = Field(None, description="Object storage URI")
+    url: Optional[str] = Field(None, description="Document download link")
+    
+    def to_api_format(self) -> Dict[str, Any]:
+        return self.model_dump(exclude_none=True)
+
+
+class KnowledgeChunk(BaseModel):
+    content: str = Field(..., description="The content of the chunk (slice)")
+    score: float = Field(..., description="Similarity score")
+    chunk_id: Optional[str] = Field(None, description="Unique identifier of the chunk")
+    doc_id: Optional[str] = Field(None, description="Unique identifier of the document")
+
+
+class KnowledgeSearchResponse(BaseModel):
+    chunks: List[KnowledgeChunk] = Field(default_factory=list, description="List of matching knowledge chunks")
+    code: int = Field(..., description="Status code")
+    msg: str = Field(..., description="Status message")
+
+
+class KnowledgeInsertResponse(BaseModel):
+    doc_ids: Optional[List[str]] = Field(None, description="List of inserted document IDs")
+    code: int = Field(..., description="Status code")
+    msg: str = Field(..., description="Status message")

local_coze/llm/__init__.py

@@ -0,0 +1,25 @@
+from .client import LLMClient
+from .models import (
+    LLMConfig,
+    ThinkingConfig,
+    CachingConfig,
+    TextContent,
+    ImageURLContent,
+    VideoURLContent,
+    ImageURLDetail,
+    VideoURLDetail,
+    MessageContent
+)
+
+__all__ = [
+    "LLMClient",
+    "LLMConfig",
+    "ThinkingConfig",
+    "CachingConfig",
+    "TextContent",
+    "ImageURLContent",
+    "VideoURLContent",
+    "ImageURLDetail",
+    "VideoURLDetail",
+    "MessageContent"
+]