local-coze 0.0.1__py3-none-any.whl

local_coze/llm/client.py

@@ -0,0 +1,317 @@
+from typing import Any, Dict, Iterator, List, Optional, Union
+
+from coze_coding_utils.runtime_ctx.context import Context, default_headers
+from cozeloop.decorator import observe
+from langchain_core.messages import (
+    AIMessage,
+    BaseMessage,
+    BaseMessageChunk,
+    HumanMessage,
+    SystemMessage,
+)
+from langchain_openai import ChatOpenAI
+
+from ..core.config import Config
+from .models import LLMConfig
+
+
+class LLMClient:
+    def __init__(
+        self,
+        config: Optional[Config] = None,
+        ctx: Optional[Context] = None,
+        custom_headers: Optional[Dict[str, str]] = None,
+        verbose: bool = False,
+    ):
+        if config is None:
+            config = Config()
+        self.config = config
+        self.ctx = ctx
+        self.custom_headers = custom_headers or {}
+        self.verbose = verbose
+        self.base_url = self.config.base_model_url
+        self.api_key = self.config.api_key
+
+    def _create_llm(
+        self,
+        llm_config: LLMConfig,
+        use_caching: bool = False,
+        previous_response_id: Optional[str] = None,
+        extra_headers: Optional[Dict[str, str]] = None,
+    ) -> ChatOpenAI:
+        extra_body = {}
+
+        if llm_config.thinking:
+            extra_body["thinking"] = {"type": llm_config.thinking}
+
+        if llm_config.caching:
+            extra_body["caching"] = {"type": llm_config.caching}
+
+        headers = {}
+
+        if self.ctx is not None:
+            ctx_headers = default_headers(self.ctx)
+            headers.update(ctx_headers)
+
+        if self.custom_headers:
+            headers.update(self.custom_headers)
+
+        if llm_config.max_tokens == 0:
+            llm_config.max_tokens = 32768
+        if llm_config.max_completion_tokens == 0:
+            llm_config.max_completion_tokens = 32768
+        if llm_config.max_tokens and llm_config.max_completion_tokens:
+            llm_config.max_tokens = None
+
+        config_headers = self.config.get_headers(extra_headers)
+        headers.update(config_headers)
+
+        llm = ChatOpenAI(
+            model=llm_config.model,
+            api_key=self.api_key,
+            base_url=self.base_url,
+            streaming=llm_config.streaming,
+            extra_body=extra_body if extra_body else None,
+            temperature=llm_config.temperature,
+            frequency_penalty=llm_config.frequency_penalty,
+            top_p=llm_config.top_p,
+            max_tokens=llm_config.max_tokens,
+            max_completion_tokens=llm_config.max_completion_tokens,
+            default_headers=headers,
+            use_responses_api=use_caching,
+            use_previous_response_id=previous_response_id is not None,
+        )
+
+        return llm
+
+    @observe(name="llm_stream")
+    def stream(
+        self,
+        messages: List[BaseMessage],
+        model: str = "doubao-seed-1-8-251228",
+        thinking: Optional[str] = "disabled",
+        caching: Optional[str] = "disabled",
+        temperature: Optional[float] = 1.0,
+        frequency_penalty: Optional[float] = 0,
+        top_p: Optional[float] = 0,
+        max_tokens: Optional[int] = None,
+        max_completion_tokens: Optional[int] = 32768,
+        previous_response_id: Optional[str] = None,
+        extra_headers: Optional[Dict[str, str]] = None,
+    ) -> Iterator[BaseMessageChunk]:
+        """
+        流式调用大语言模型，逐块返回生成的内容
+
+        Args:
+            messages: 消息列表，使用 LangChain 消息格式（必需）
+                - SystemMessage: 系统提示词，定义 AI 角色和行为
+                - HumanMessage: 用户消息
+                - AIMessage: AI 回复，用于多轮对话
+
+            model: 模型ID，默认 "doubao-seed-1-8-251228"
+                可选模型：
+                - "doubao-seed-1-8-251228": 最新模型，更高性能
+                - "doubao-seed-1-6-251015": 平衡性能
+                - "doubao-seed-1-6-flash-250615": 快速模型
+                - "doubao-seed-1-6-thinking-250715": 思考模型
+
+            thinking: 思考模式，默认 "disabled"
+                - "enabled": 启用深度思考，适合复杂推理任务
+                - "disabled": 禁用，适合快速响应
+
+            caching: 缓存模式，默认 "disabled"
+                - "enabled": 启用缓存，加速重复上下文的响应
+                - "disabled": 禁用
+
+            temperature: 温度参数，控制输出随机性，范围 0-2，默认 1.0
+                - 0.0-0.3: 确定性输出，适合代码生成、数据分析
+                - 0.7-0.9: 平衡创造性，适合通用对话
+                - 1.0-2.0: 高创造性，适合创意写作、头脑风暴
+
+            frequency_penalty: 频率惩罚，减少重复内容，范围 -2 到 2，默认 0
+                正值减少重复，负值增加重复
+
+            top_p: 核采样参数，控制输出多样性，范围 0-1，默认 0
+                值越小输出越确定，值越大输出越多样
+
+            max_tokens: 最大输出 token 数，默认 None（不限制）
+                用于限制输出长度或控制成本
+
+            max_completion_tokens: 最大完成 token 数，默认 None
+                更精确的长度控制
+
+            previous_response_id: 上一次响应ID，用于缓存场景，默认 None
+
+            extra_headers: 额外的 HTTP 请求头，默认 None
+
+        Returns:
+            Iterator[BaseMessageChunk]: 流式返回的消息块，每个块包含：
+                - content: 文本内容
+                - response_metadata: 响应元数据
+
+        Example:
+            >>> from langchain_core.messages import HumanMessage
+            >>> client = LLMClient()
+            >>> messages = [HumanMessage(content="你好")]
+            >>>
+            >>> # 最简单用法
+            >>> for chunk in client.stream(messages):
+            ...     if chunk.content:
+            ...         print(chunk.content, end="")
+            >>>
+            >>> # 调整温度
+            >>> for chunk in client.stream(messages, temperature=0.7):
+            ...     if chunk.content:
+            ...         print(chunk.content, end="")
+            >>>
+            >>> # 启用思考模式
+            >>> for chunk in client.stream(messages, thinking="enabled"):
+            ...     if chunk.content:
+            ...         print(chunk.content, end="")
+        """
+        llm_config = LLMConfig(
+            model=model,
+            thinking=thinking,
+            caching=caching,
+            temperature=temperature,
+            frequency_penalty=frequency_penalty,
+            top_p=top_p,
+            max_tokens=max_tokens,
+            max_completion_tokens=max_completion_tokens,
+            streaming=True,
+        )
+
+        use_caching = caching == "enabled" or previous_response_id is not None
+
+        if previous_response_id:
+            for i in range(len(messages) - 1, -1, -1):
+                msg = messages[i]
+                if isinstance(msg, AIMessage):
+                    msg.response_metadata["id"] = previous_response_id
+                    break
+
+        llm = self._create_llm(
+            llm_config,
+            use_caching=use_caching,
+            previous_response_id=previous_response_id,
+            extra_headers=extra_headers,
+        )
+
+        for chunk in llm.stream(messages):
+            yield chunk
+
+    @observe(name="llm_invoke")
+    def invoke(
+        self,
+        messages: List[BaseMessage],
+        model: str = "doubao-seed-1-8-251228",
+        thinking: Optional[str] = "disabled",
+        caching: Optional[str] = "disabled",
+        temperature: Optional[float] = 1.0,
+        frequency_penalty: Optional[float] = 0,
+        top_p: Optional[float] = 0,
+        max_tokens: Optional[int] = None,
+        max_completion_tokens: Optional[int] = 32768,
+        previous_response_id: Optional[str] = None,
+        extra_headers: Optional[Dict[str, str]] = None,
+    ) -> AIMessage:
+        """
+        非流式调用大语言模型，返回完整的响应
+
+        内部通过流式调用实现，自动组装完整响应后返回。
+
+        Args:
+            messages: 消息列表，使用 LangChain 消息格式（必需）
+                - SystemMessage: 系统提示词，定义 AI 角色和行为
+                - HumanMessage: 用户消息
+                - AIMessage: AI 回复，用于多轮对话
+
+            model: 模型ID，默认 "doubao-seed-1-8-251228"
+                可选模型：
+                - "doubao-seed-1-8-251228": 最新模型，更高性能
+                - "doubao-seed-1-6-251015": 平衡性能
+                - "doubao-seed-1-6-flash-250615": 快速模型
+                - "doubao-seed-1-6-thinking-250715": 思考模型
+
+            thinking: 思考模式，默认 "disabled"
+                - "enabled": 启用深度思考，适合复杂推理任务
+                - "disabled": 禁用，适合快速响应
+
+            caching: 缓存模式，默认 "disabled"
+                - "enabled": 启用缓存，加速重复上下文的响应
+                - "disabled": 禁用
+
+            temperature: 温度参数，控制输出随机性，范围 0-2，默认 1.0
+                - 0.0-0.3: 确定性输出，适合代码生成、数据分析
+                - 0.7-0.9: 平衡创造性，适合通用对话
+                - 1.0-2.0: 高创造性，适合创意写作、头脑风暴
+
+            frequency_penalty: 频率惩罚，减少重复内容，范围 -2 到 2，默认 0
+                正值减少重复，负值增加重复
+
+            top_p: 核采样参数，控制输出多样性，范围 0-1，默认 0
+                值越小输出越确定，值越大输出越多样
+
+            max_tokens: 最大输出 token 数，默认 None（不限制）
+                用于限制输出长度或控制成本
+
+            max_completion_tokens: 最大完成 token 数，默认 None
+                更精确的长度控制
+
+            previous_response_id: 上一次响应ID，用于缓存场景，默认 None
+
+            extra_headers: 额外的 HTTP 请求头，默认 None
+
+        Returns:
+            AIMessage: 完整的响应消息，包含：
+                - content: 完整的文本内容
+                - response_metadata: 响应元数据（模型信息、token 使用量等）
+
+        Example:
+            >>> from langchain_core.messages import SystemMessage, HumanMessage
+            >>> client = LLMClient()
+            >>>
+            >>> # 最简单用法
+            >>> messages = [HumanMessage(content="你好")]
+            >>> response = client.invoke(messages)
+            >>> print(response.content)
+            >>>
+            >>> # 带系统提示词
+            >>> messages = [
+            ...     SystemMessage(content="你是一个 Python 专家"),
+            ...     HumanMessage(content="什么是装饰器？")
+            ... ]
+            >>> response = client.invoke(messages)
+            >>> print(response.content)
+            >>>
+            >>> # 调整参数
+            >>> response = client.invoke(
+            ...     messages=messages,
+            ...     temperature=0.7,
+            ...     max_tokens=500
+            ... )
+            >>> print(response.content)
+            >>> print(f"Token 使用: {response.response_metadata}")
+        """
+        full_content = ""
+        response_metadata = {}
+
+        for chunk in self.stream(
+            messages=messages,
+            model=model,
+            thinking=thinking,
+            caching=caching,
+            temperature=temperature,
+            frequency_penalty=frequency_penalty,
+            top_p=top_p,
+            max_tokens=max_tokens,
+            max_completion_tokens=max_completion_tokens,
+            previous_response_id=previous_response_id,
+            extra_headers=extra_headers,
+        ):
+            if chunk.content:
+                full_content += str(chunk.content)
+            if chunk.response_metadata:
+                response_metadata.update(chunk.response_metadata)
+
+        return AIMessage(content=full_content, response_metadata=response_metadata)

local_coze/llm/models.py

@@ -0,0 +1,48 @@
+from typing import Optional, Literal, Dict, Any, List, Union
+from pydantic import BaseModel, Field
+
+
+class ThinkingConfig(BaseModel):
+    type: Literal["enabled", "disabled"] = Field("disabled", description="是否开启深度思考能力")
+
+
+class CachingConfig(BaseModel):
+    type: Literal["enabled", "disabled"] = Field("disabled", description="是否开启模型缓存")
+
+
+class LLMConfig(BaseModel):
+    model: str = Field("doubao-seed-1-8-251228", description="模型ID")
+    thinking: Optional[Literal["enabled", "disabled"]] = Field("disabled", description="是否开启深度思考能力")
+    caching: Optional[Literal["enabled", "disabled"]] = Field("disabled", description="是否开启模型缓存")
+    temperature: Optional[float] = Field(1.0, ge=0, le=2, description="控制模型输出的随机性")
+    frequency_penalty: Optional[float] = Field(0, ge=-2, le=2, description="重复语句惩罚")
+    top_p: Optional[float] = Field(0, ge=0, le=1, description="控制模型输出的多样性")
+    max_tokens: Optional[int] = Field(None, description="控制模型输出的最大 tokens 数")
+    max_completion_tokens: Optional[int] = Field(None, description="控制模型输出的最大 completion tokens 数")
+    streaming: bool = Field(True, description="是否使用流式输出")
+
+
+class TextContent(BaseModel):
+    type: Literal["text"] = "text"
+    text: str
+
+
+class ImageURLDetail(BaseModel):
+    url: str
+
+
+class ImageURLContent(BaseModel):
+    type: Literal["image_url"] = "image_url"
+    image_url: ImageURLDetail
+
+
+class VideoURLDetail(BaseModel):
+    url: str
+
+
+class VideoURLContent(BaseModel):
+    type: Literal["video_url"] = "video_url"
+    video_url: VideoURLDetail
+
+
+MessageContent = Union[str, List[Union[TextContent, ImageURLContent, VideoURLContent]]]

local_coze/memory/__init__.py

@@ -0,0 +1,14 @@
+"""
+Memory 模块
+提供 LangGraph 检查点管理
+"""
+
+from .client import (
+    MemoryManager,
+    get_memory_saver,
+)
+
+__all__ = [
+    "MemoryManager",
+    "get_memory_saver",
+]

local_coze/memory/client.py

@@ -0,0 +1,176 @@
+"""
+Memory 模块
+提供 LangGraph 检查点管理，支持 PostgreSQL 持久化和内存兜底
+"""
+
+import os
+import time
+import logging
+from typing import Optional, Union
+
+import psycopg
+from psycopg_pool import AsyncConnectionPool
+from langgraph.checkpoint.postgres import PostgresSaver
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.checkpoint.base import BaseCheckpointSaver
+
+logger = logging.getLogger(__name__)
+
+# 数据库连接超时时间（秒），每次尝试 15 秒，共尝试 2 次
+DB_CONNECTION_TIMEOUT = 15
+DB_MAX_RETRIES = 2
+
+
+def _load_env() -> None:
+    """加载环境变量"""
+    try:
+        from dotenv import load_dotenv
+        load_dotenv()
+    except ImportError:
+        pass
+
+    try:
+        from coze_workload_identity import Client
+        client = Client()
+        env_vars = client.get_project_env_vars()
+        client.close()
+        for env_var in env_vars:
+            if env_var.key not in os.environ:
+                os.environ[env_var.key] = env_var.value
+    except Exception as e:
+        logger.debug(f"coze_workload_identity not available: {e}")
+
+
+def _get_db_url() -> Optional[str]:
+    """安全获取 db_url，失败时返回 None"""
+    _load_env()
+
+    url = os.getenv("PGDATABASE_URL")
+    if url and url.strip():
+        return url
+
+    logger.warning("PGDATABASE_URL is not set, will fallback to MemorySaver")
+    return None
+
+
+class MemoryManager:
+    """Memory Manager 单例类"""
+
+    _instance: Optional['MemoryManager'] = None
+    _checkpointer: Optional[Union[AsyncPostgresSaver, MemorySaver]] = None
+    _pool: Optional[AsyncConnectionPool] = None
+    _setup_done: bool = False
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def _connect_with_retry(self, db_url: str) -> Optional[psycopg.Connection]:
+        """带重试的数据库连接，每次 15 秒超时，共尝试 2 次"""
+        last_error = None
+        for attempt in range(1, DB_MAX_RETRIES + 1):
+            try:
+                logger.info(f"Attempting database connection (attempt {attempt}/{DB_MAX_RETRIES})")
+                conn = psycopg.connect(db_url, autocommit=True, connect_timeout=DB_CONNECTION_TIMEOUT)
+                logger.info(f"Database connection established on attempt {attempt}")
+                return conn
+            except Exception as e:
+                last_error = e
+                logger.warning(f"Database connection attempt {attempt} failed: {e}")
+                if attempt < DB_MAX_RETRIES:
+                    time.sleep(1)  # 重试前短暂等待
+        logger.error(f"All {DB_MAX_RETRIES} database connection attempts failed, last error: {last_error}")
+        return None
+
+    def _setup_schema_and_tables(self, db_url: str) -> bool:
+        """同步创建 schema 和表（只执行一次），返回是否成功"""
+        if self._setup_done:
+            return True
+
+        conn = self._connect_with_retry(db_url)
+        if conn is None:
+            return False
+
+        try:
+            with conn.cursor() as cur:
+                cur.execute("CREATE SCHEMA IF NOT EXISTS memory")
+            conn.execute("SET search_path TO memory")
+            PostgresSaver(conn).setup()
+            self._setup_done = True
+            logger.info("Memory schema and tables created")
+            return True
+        except Exception as e:
+            logger.warning(f"Failed to setup schema/tables: {e}")
+            return False
+        finally:
+            conn.close()
+
+    def _create_fallback_checkpointer(self) -> MemorySaver:
+        """创建内存兜底 checkpointer"""
+        self._checkpointer = MemorySaver()
+        logger.warning("Using MemorySaver as fallback checkpointer (data will not persist across restarts)")
+        return self._checkpointer
+
+    def get_checkpointer(self) -> BaseCheckpointSaver:
+        """获取 checkpointer，优先使用 PostgresSaver，失败时退化为 MemorySaver"""
+        if self._checkpointer is not None:
+            return self._checkpointer
+
+        # 1. 尝试获取 db_url
+        db_url = _get_db_url()
+        if not db_url:
+            return self._create_fallback_checkpointer()
+
+        # 2. 尝试连接数据库并创建 schema/表（带重试）
+        if not self._setup_schema_and_tables(db_url):
+            return self._create_fallback_checkpointer()
+
+        # 3. 连接字符串加上 search_path
+        if "?" in db_url:
+            db_url = f"{db_url}&options=-csearch_path%3Dmemory"
+        else:
+            db_url = f"{db_url}?options=-csearch_path%3Dmemory"
+
+        # 4. 尝试创建连接池和 checkpointer
+        try:
+            self._pool = AsyncConnectionPool(
+                conninfo=db_url,
+                timeout=DB_CONNECTION_TIMEOUT,
+                min_size=1,
+                max_idle=300,
+            )
+            self._checkpointer = AsyncPostgresSaver(self._pool)
+            logger.info("AsyncPostgresSaver initialized successfully")
+        except Exception as e:
+            logger.warning(f"Failed to create AsyncPostgresSaver: {e}, will fallback to MemorySaver")
+            return self._create_fallback_checkpointer()
+
+        return self._checkpointer
+
+
+_memory_manager: Optional[MemoryManager] = None
+
+
+def get_memory_saver() -> BaseCheckpointSaver:
+    """
+    获取 checkpointer，优先使用 PostgresSaver，db_url 不可用或连接失败时退化为 MemorySaver
+
+    Returns:
+        BaseCheckpointSaver: LangGraph 检查点保存器
+
+    Example:
+        >>> from local_coze.memory import get_memory_saver
+        >>> checkpointer = get_memory_saver()
+    """
+    global _memory_manager
+    if _memory_manager is None:
+        _memory_manager = MemoryManager()
+    return _memory_manager.get_checkpointer()
+
+
+__all__ = [
+    "MemoryManager",
+    "get_memory_saver",
+]

local_coze/s3/__init__.py

@@ -0,0 +1,12 @@
+"""
+S3 兼容存储模块
+提供对象存储的上传、下载、删除等功能
+"""
+
+from .client import S3SyncStorage
+from .models import ListFilesResult
+
+__all__ = [
+    "S3SyncStorage",
+    "ListFilesResult",
+]