local-coze 0.0.1__py3-none-any.whl

local_coze/search/client.py

@@ -0,0 +1,183 @@
+from typing import Dict, List, Optional
+
+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 ImageItem, SearchFilter, SearchRequest, SearchResponse, WebItem
+
+
+def _convert_to_api_format(data: dict) -> dict:
+    result = {}
+    for key, value in data.items():
+        if value is None:
+            continue
+        if key == "query":
+            result["Query"] = value
+        elif key == "search_type":
+            result["SearchType"] = value
+        elif key == "count":
+            result["Count"] = value
+        elif key == "need_summary":
+            result["NeedSummary"] = value
+        elif key == "time_range":
+            result["TimeRange"] = value
+        elif key == "filter" and value:
+            result["Filter"] = {
+                "NeedContent": value.get("need_content", False),
+                "NeedUrl": value.get("need_url", False),
+                "Sites": value.get("sites"),
+                "BlockHosts": value.get("block_hosts"),
+            }
+    return result
+
+
+def _convert_from_api_format(item: dict) -> dict:
+    result = {}
+    for key, value in item.items():
+        if key == "Id":
+            result["id"] = value
+        elif key == "SortId":
+            result["sort_id"] = value
+        elif key == "Title":
+            result["title"] = value
+        elif key == "SiteName":
+            result["site_name"] = value
+        elif key == "Url":
+            result["url"] = value
+        elif key == "Snippet":
+            result["snippet"] = value
+        elif key == "Summary":
+            result["summary"] = value
+        elif key == "Content":
+            result["content"] = value
+        elif key == "PublishTime":
+            result["publish_time"] = value
+        elif key == "LogoUrl":
+            result["logo_url"] = value
+        elif key == "RankScore":
+            result["rank_score"] = value
+        elif key == "AuthInfoDes":
+            result["auth_info_des"] = value
+        elif key == "AuthInfoLevel":
+            result["auth_info_level"] = value
+        elif key == "Image":
+            result["image"] = {
+                "url": value.get("Url"),
+                "width": value.get("Width"),
+                "height": value.get("Height"),
+                "shape": value.get("Shape"),
+            }
+    return result
+
+
+class SearchClient(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
+
+    @observe(name="web_search")
+    def search(
+        self,
+        query: str,
+        search_type: str = "web",
+        count: Optional[int] = 10,
+        need_content: Optional[bool] = False,
+        need_url: Optional[bool] = False,
+        sites: Optional[str] = None,
+        block_hosts: Optional[str] = None,
+        need_summary: Optional[bool] = True,
+        time_range: Optional[str] = None,
+    ) -> SearchResponse:
+        search_filter = SearchFilter(
+            need_content=need_content,
+            need_url=need_url,
+            sites=sites,
+            block_hosts=block_hosts,
+        )
+
+        request = SearchRequest(
+            query=query,
+            search_type=search_type,
+            count=count,
+            filter=search_filter,
+            need_summary=need_summary,
+            time_range=time_range,
+        )
+
+        api_request = _convert_to_api_format(request.model_dump(exclude_none=True))
+
+        response = self._request(
+            method="POST",
+            url=f"{self.base_url}/api/search_api/web_search",
+            json=api_request,
+        )
+
+        response_metadata = response.get("ResponseMetadata", {})
+        if response_metadata.get("Error"):
+            raise APIError(f"Search failed: {response_metadata.get('Error')}")
+
+        result = response.get("Result", {})
+
+        web_items = []
+        if result.get("WebResults"):
+            web_items = [
+                WebItem(**_convert_from_api_format(item))
+                for item in result.get("WebResults", [])
+            ]
+
+        image_items = []
+        if result.get("ImageResults"):
+            image_items = [
+                ImageItem(**_convert_from_api_format(item))
+                for item in result.get("ImageResults", [])
+            ]
+
+        summary = None
+        if result.get("Choices"):
+            summary = (
+                result.get("Choices", [{}])[0].get("Message", {}).get("Content", "")
+            )
+
+        return SearchResponse(
+            web_items=web_items, image_items=image_items, summary=summary
+        )
+
+    @observe(name="web_search_simple")
+    def web_search(
+        self,
+        query: str,
+        count: Optional[int] = 10,
+        need_summary: Optional[bool] = True,
+    ) -> SearchResponse:
+        return self.search(
+            query=query, search_type="web", count=count, need_summary=need_summary
+        )
+
+    @observe(name="web_search_with_summary")
+    def web_search_with_summary(
+        self,
+        query: str,
+        count: Optional[int] = 10,
+    ) -> SearchResponse:
+        return self.search(
+            query=query, search_type="web_summary", count=count, need_summary=True
+        )
+
+    @observe(name="image_search")
+    def image_search(
+        self,
+        query: str,
+        count: Optional[int] = 10,
+    ) -> SearchResponse:
+        return self.search(
+            query=query, search_type="image", count=count, need_summary=False
+        )

local_coze/search/models.py

@@ -0,0 +1,57 @@
+from typing import Optional, List, Literal
+from pydantic import BaseModel, Field
+
+
+class ImageInfo(BaseModel):
+    url: str = Field(..., description="图片链接")
+    width: Optional[int] = Field(None, description="宽")
+    height: Optional[int] = Field(None, description="高")
+    shape: str = Field(..., description="横长方形/竖长方形/方形")
+
+
+class WebItem(BaseModel):
+    id: str = Field(..., description="结果Id")
+    sort_id: int = Field(..., description="排序Id")
+    title: str = Field(..., description="标题")
+    site_name: Optional[str] = Field(None, description="站点名")
+    url: Optional[str] = Field(None, description="落地页")
+    snippet: str = Field(..., description="普通摘要")
+    summary: Optional[str] = Field(None, description="精准摘要")
+    content: Optional[str] = Field(None, description="正文")
+    publish_time: Optional[str] = Field(None, description="发布时间")
+    logo_url: Optional[str] = Field(None, description="落地页IconUrl链接")
+    rank_score: Optional[float] = Field(None, description="得分")
+    auth_info_des: str = Field(..., description="权威度描述")
+    auth_info_level: int = Field(..., description="权威度评级")
+
+
+class ImageItem(BaseModel):
+    id: str = Field(..., description="结果Id")
+    sort_id: int = Field(..., description="排序Id")
+    title: Optional[str] = Field(None, description="标题")
+    site_name: Optional[str] = Field(None, description="站点名")
+    url: Optional[str] = Field(None, description="落地页")
+    publish_time: Optional[str] = Field(None, description="发布时间")
+    image: ImageInfo = Field(..., description="图片详情")
+
+
+class SearchFilter(BaseModel):
+    need_content: Optional[bool] = Field(False, description="是否仅返回有正文的结果")
+    need_url: Optional[bool] = Field(False, description="是否仅返回原文链接的结果")
+    sites: Optional[str] = Field(None, description="指定搜索的Site范围")
+    block_hosts: Optional[str] = Field(None, description="指定屏蔽的搜索Site")
+
+
+class SearchRequest(BaseModel):
+    query: str = Field(..., description="用户搜索query")
+    search_type: Literal["web", "web_summary", "image"] = Field("web", description="搜索类型")
+    count: Optional[int] = Field(10, description="返回条数")
+    filter: Optional[SearchFilter] = None
+    need_summary: Optional[bool] = Field(True, description="是否需要精准摘要")
+    time_range: Optional[str] = Field(None, description="指定搜索的发文时间")
+
+
+class SearchResponse(BaseModel):
+    web_items: List[WebItem] = Field(default_factory=list, description="Web搜索结果")
+    image_items: List[ImageItem] = Field(default_factory=list, description="图片搜索结果")
+    summary: Optional[str] = Field(None, description="搜索结果摘要")

local_coze/video/__init__.py

@@ -0,0 +1,17 @@
+from .client import VideoGenerationClient
+from .models import (
+    VideoGenerationRequest,
+    VideoGenerationTask,
+    TextContent,
+    ImageURLContent,
+    ImageURL
+)
+
+__all__ = [
+    "VideoGenerationClient",
+    "VideoGenerationRequest",
+    "VideoGenerationTask",
+    "TextContent",
+    "ImageURLContent",
+    "ImageURL"
+]

local_coze/video/client.py

@@ -0,0 +1,347 @@
+import asyncio
+import logging
+import time
+from typing import Dict, List, Optional, Tuple, 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 (
+    ImageURLContent,
+    TextContent,
+    VideoGenerationRequest,
+    VideoGenerationTask,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class VideoGenerationClient(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
+
+    def _create_task(
+        self,
+        model: str,
+        content: List[Union[TextContent, ImageURLContent]],
+        resolution: Optional[str] = "720p",
+        ratio: Optional[str] = "16:9",
+        duration: Optional[int] = 5,
+        watermark: Optional[bool] = True,
+        seed: Optional[int] = None,
+        camerafixed: Optional[bool] = False,
+        generate_audio: Optional[bool] = True,
+    ) -> str:
+        request = VideoGenerationRequest(
+            model=model,
+            content=content,
+            resolution=resolution,
+            ratio=ratio,
+            duration=duration,
+            watermark=watermark,
+            seed=seed,
+            camerafixed=camerafixed,
+            generate_audio=generate_audio,
+        )
+
+        response = self._request(
+            method="POST",
+            url=f"{self.base_url}/api/v3/contents/generations/tasks",
+            json=request.model_dump(exclude_none=True),
+        )
+
+        return response.get("id")
+
+    def _get_task_status(self, task_id: str) -> VideoGenerationTask:
+        response = self._request(
+            method="GET",
+            url=f"{self.base_url}/api/v3/contents/generations/tasks/{task_id}",
+        )
+
+        return VideoGenerationTask(**response)
+
+    @observe(name="video_generation")
+    def video_generation(
+        self,
+        content_items: List[Union[TextContent, ImageURLContent]],
+        callback_url: Optional[str] = None,
+        return_last_frame: Optional[bool] = False,
+        model: str = "doubao-seedance-1-5-pro-251215",
+        max_wait_time: int = 900,
+        resolution: Optional[str] = "720p",
+        ratio: Optional[str] = "16:9",
+        duration: Optional[int] = 5,
+        watermark: Optional[bool] = True,
+        seed: Optional[int] = None,
+        camerafixed: Optional[bool] = False,
+        generate_audio: Optional[bool] = True,
+    ) -> Tuple[Optional[str], Dict, str]:
+        """
+        视频生成函数（同步）。调用生视频模型，非大语言模型，仅用于生成视频
+
+        Args:
+            content_items: 输入给模型，生成视频的信息，支持文本信息和图片信息。
+            callback_url: 可选，填写本次生成任务结果的回调通知地址。当视频生成任务有状态变化时，方舟将向此地址推送 POST 请求。
+            return_last_frame: 可选，返回生成视频的尾帧图像。设置为 true 后，可获取视频的尾帧图像。
+            model: 模型名称，默认使用 doubao-seedance-1-5-pro-251215
+            max_wait_time: 最大等待时间（秒），默认 900 秒
+            resolution: 视频分辨率，可选 "480p" 或 "720p"，默认 "720p"
+            ratio: 视频比例，可选 "16:9", "9:16", "1:1", "4:3", "3:4", "21:9", "adaptive"，默认 "16:9"
+            duration: 视频时长（秒），范围 4 到 12，默认 5，若是 -1，则由模型决定长度
+            watermark: 是否添加水印，默认 True
+            seed: 随机种子，用于复现生成结果
+            camerafixed: 是否固定摄像机位置，默认 False
+            generate_audio: 是否生成音频，默认 True
+
+        Returns:
+            Tuple[Optional[str], Dict, str]: (视频URL, 完整响应数据字典, 尾帧图像URL)
+                完整响应数据示例：
+                {
+                  "id": "cgt-2025******-****",
+                  "model": "doubao-seedance-1-5-pro-251215",
+                  "status": "succeeded",
+                  "content": {
+                    "video_url": "https://ark-content-generation-cn-beijing.tos-cn-beijing.volces.com/..."
+                  },
+                  "seed": 10,
+                  "resolution": "720p",
+                  "ratio": "16:9",
+                  "duration": 5,
+                  "framespersecond": 24,
+                  "usage": {
+                    "completion_tokens": 108900,
+                    "total_tokens": 108900
+                  },
+                  "created_at": 1743414619,
+                  "updated_at": 1743414673
+                }
+
+        Raises:
+            APIError: 当视频生成失败、超时或出现其他错误时
+        """
+        poll_interval = 5
+
+        request_data = {
+            "model": model,
+            "content": [item.model_dump() for item in content_items],
+        }
+
+        if callback_url:
+            request_data["callback_url"] = callback_url
+        if return_last_frame:
+            request_data["return_last_frame"] = return_last_frame
+
+        if resolution is not None:
+            request_data["resolution"] = resolution
+        if ratio is not None:
+            request_data["ratio"] = ratio
+        if duration is not None:
+            if duration < 4 or duration > 12:
+                # 兜底策略
+                duration = -1
+            request_data["duration"] = duration
+        if watermark is not None:
+            request_data["watermark"] = watermark
+        if seed is not None:
+            request_data["seed"] = seed
+        if camerafixed is not None:
+            request_data["camerafixed"] = camerafixed
+        if generate_audio is not None:
+            request_data["generate_audio"] = generate_audio
+
+        task_id = None
+        retry_count = 0
+        max_retries = 3
+
+        while retry_count < max_retries:
+            try:
+                response = self._request(
+                    method="POST",
+                    url=f"{self.base_url}/api/v3/contents/generations/tasks",
+                    json=request_data,
+                )
+
+                task_id = response.get("id")
+                if not task_id:
+                    raise APIError("创建视频生成任务失败：响应中缺少任务ID")
+
+                logger.info(f"视频生成任务创建成功，任务ID: {task_id}")
+                break
+
+            except APIError as e:
+                retry_count += 1
+                error_msg = str(e)
+                logger.error(
+                    f"创建视频生成任务失败（尝试 {retry_count}/{max_retries}）: {error_msg}"
+                )
+
+                if retry_count >= max_retries:
+                    raise APIError(
+                        f"创建视频生成任务失败，已重试{max_retries}次: {error_msg}"
+                    )
+
+                if "rate limit" in error_msg.lower() or "429" in error_msg:
+                    wait_time = min(2**retry_count, 10)
+                    logger.warning(f"遇到速率限制，等待 {wait_time} 秒后重试...")
+                    time.sleep(wait_time)
+                elif "timeout" in error_msg.lower():
+                    logger.warning(f"请求超时，等待 2 秒后重试...")
+                    time.sleep(2)
+                else:
+                    raise
+
+        start_time = time.time()
+        poll_count = 0
+
+        while time.time() - start_time < max_wait_time:
+            poll_count += 1
+
+            try:
+                response = self._request(
+                    method="GET",
+                    url=f"{self.base_url}/api/v3/contents/generations/tasks/{task_id}",
+                )
+
+                status = response.get("status")
+                logger.debug(f"任务 {task_id} 状态检查 #{poll_count}: {status}")
+
+                if status == "succeeded":
+                    video_url = response.get("content", {}).get("video_url")
+                    last_frame_url = response.get("content", {}).get(
+                        "last_frame_url", ""
+                    )
+
+                    if not video_url:
+                        raise APIError(
+                            f"视频生成成功但响应中缺少视频URL，任务ID: {task_id}"
+                        )
+
+                    logger.info(
+                        f"视频生成成功，任务ID: {task_id}, 视频URL: {video_url}"
+                    )
+                    return video_url, response, last_frame_url
+
+                elif status == "failed":
+                    error_message = response.get("error_message", "未知错误")
+                    logger.error(
+                        f"视频生成失败，任务ID: {task_id}, 错误: {error_message}"
+                    )
+                    raise APIError(f"视频生成失败: {error_message}")
+
+                elif status == "cancelled":
+                    logger.warning(f"视频生成任务被取消，任务ID: {task_id}")
+                    return None, response, ""
+
+                elif status in ["queued", "running"]:
+                    time.sleep(poll_interval)
+                    continue
+
+                else:
+                    logger.warning(f"未知的任务状态: {status}，任务ID: {task_id}")
+                    time.sleep(poll_interval)
+                    continue
+
+            except APIError as e:
+                error_msg = str(e)
+                logger.error(f"查询任务状态失败，任务ID: {task_id}, 错误: {error_msg}")
+
+                if "not found" in error_msg.lower() or "404" in error_msg:
+                    raise APIError(f"任务不存在或已过期，任务ID: {task_id}")
+
+                if time.time() - start_time >= max_wait_time:
+                    raise
+
+                time.sleep(poll_interval)
+                continue
+
+        elapsed_time = int(time.time() - start_time)
+        logger.error(f"视频生成超时，任务ID: {task_id}, 已等待: {elapsed_time}秒")
+        raise APIError(f"视频生成超时，已等待 {elapsed_time} 秒，任务ID: {task_id}")
+
+    async def video_generation_async(
+        self,
+        content_items: List[Union[TextContent, ImageURLContent]],
+        callback_url: Optional[str] = None,
+        return_last_frame: Optional[bool] = False,
+        model: str = "doubao-seedance-1-5-pro-251215",
+        max_wait_time: int = 900,
+        resolution: Optional[str] = "720p",
+        ratio: Optional[str] = "16:9",
+        duration: Optional[int] = 5,
+        watermark: Optional[bool] = True,
+        seed: Optional[int] = None,
+        camerafixed: Optional[bool] = False,
+        generate_audio: Optional[bool] = True,
+    ) -> Tuple[Optional[str], Dict, str]:
+        """
+        视频生成函数（异步）。调用生视频模型，非大语言模型，仅用于生成视频
+
+        适用于批量生成视频的场景，可以并发执行多个视频生成任务。
+
+        Example:
+            ```python
+            import asyncio
+
+            async def batch_generate():
+                prompts = ["小猫玩球", "日落海滩", "城市夜景"]
+                tasks = [
+                    client.video_generation_async(
+                        content_items=[TextContent(text=prompt)]
+                    )
+                    for prompt in prompts
+                ]
+                results = await asyncio.gather(*tasks)
+                return results
+
+            results = asyncio.run(batch_generate())
+            ```
+
+        Args:
+            content_items: 输入给模型，生成视频的信息，支持文本信息和图片信息。
+            callback_url: 可选，填写本次生成任务结果的回调通知地址。
+            return_last_frame: 可选，返回生成视频的尾帧图像。
+            model: 模型名称，默认使用 doubao-seedance-1-5-pro-251215
+            max_wait_time: 最大等待时间（秒），默认 900 秒
+            resolution: 视频分辨率，可选 "480p" 或 "720p"，默认 "720p"
+            ratio: 视频比例，可选 "16:9", "9:16", "1:1", "4:3", "3:4", "21:9", "adaptive"，默认 "16:9"
+            duration: 视频时长（秒），范围 -1 到 12，默认 5
+            watermark: 是否添加水印，默认 True
+            seed: 随机种子，用于复现生成结果
+            camerafixed: 是否固定摄像机位置，默认 False
+            generate_audio: 是否生成音频，默认 True
+
+        Returns:
+            Tuple[Optional[str], Dict, str]: (视频URL, 完整响应数据字典, 尾帧图像URL)
+
+        Raises:
+            APIError: 当视频生成失败、超时或出现其他错误时
+        """
+        loop = asyncio.get_event_loop()
+        result = await loop.run_in_executor(
+            None,
+            lambda: self.video_generation(
+                content_items=content_items,
+                callback_url=callback_url,
+                return_last_frame=return_last_frame,
+                model=model,
+                max_wait_time=max_wait_time,
+                resolution=resolution,
+                ratio=ratio,
+                duration=duration,
+                watermark=watermark,
+                seed=seed,
+                camerafixed=camerafixed,
+                generate_audio=generate_audio,
+            ),
+        )
+        return result

local_coze/video/models.py

@@ -0,0 +1,39 @@
+from typing import List, Literal, Optional, Union
+
+from pydantic import BaseModel, Field
+
+
+class ImageURL(BaseModel):
+    url: str
+
+
+class ImageURLContent(BaseModel):
+    type: Literal["image_url"] = "image_url"
+    image_url: ImageURL
+    role: Optional[Literal["first_frame", "last_frame", "reference_image"]] = None
+
+
+class TextContent(BaseModel):
+    type: Literal["text"] = "text"
+    text: str
+
+
+class VideoGenerationRequest(BaseModel):
+    model: str
+    content: List[Union[TextContent, ImageURLContent]]
+    resolution: Optional[Literal["480p", "720p", "1080p"]] = "720p"
+    ratio: Optional[
+        Literal["16:9", "9:16", "1:1", "4:3", "3:4", "21:9", "adaptive"]
+    ] = "16:9"
+    duration: Optional[int] = Field(default=5, ge=-1, le=12)
+    watermark: Optional[bool] = True
+    seed: Optional[int] = None
+    camerafixed: Optional[bool] = False
+    generate_audio: Optional[bool] = True
+
+
+class VideoGenerationTask(BaseModel):
+    id: str
+    status: Literal["processing", "completed", "failed"]
+    video_url: Optional[str] = None
+    error_message: Optional[str] = None

local_coze/video_edit/__init__.py

@@ -0,0 +1,23 @@
+from .frame_extractor import FrameExtractorClient
+from .video_edit import VideoEditClient
+from .models import (
+    FrameExtractorResponse,
+    FrameChunk,
+    VideoEditResponse,
+    SubtitleConfig,
+    FontPosConfig,
+    TextItem,
+    OutputSync,
+)
+
+__all__ = [
+    "FrameExtractorClient",
+    "VideoEditClient",
+    "FrameExtractorResponse",
+    "FrameChunk",
+    "VideoEditResponse",
+    "SubtitleConfig",
+    "FontPosConfig",
+    "TextItem",
+    "OutputSync",
+]