[feat][prompt] prompt as a service (ptaas) (#28)

ptaas

Co-Authored-By: Coda <coda@bytedance.com>

Author: kasarolzzw

.gitignore

@@ -34,4 +34,5 @@ venv.bak/
 /.idea
 /.vscode
 /output
-dist/
+dist/
+.coda/

cozeloop/__init__.py

@@ -18,6 +18,8 @@
     close,
     get_prompt,
     prompt_format,
+    execute_prompt,
+    aexecute_prompt,
     start_span,
     get_span_from_context,
     get_span_from_header,
@@ -30,4 +32,4 @@
     ENV_JWT_OAUTH_PUBLIC_KEY_ID
 )
 
-from .span import SpanContext, Span
+from .span import SpanContext, Span

cozeloop/_client.py

@@ -7,13 +7,14 @@
 import os
 import threading
 from datetime import datetime
-from typing import Dict, Any, List, Optional, Callable
+from typing import Dict, Any, List, Optional, Callable, Union
 
 import httpx
 
 from cozeloop.client import Client
 from cozeloop._noop import NOOP_SPAN, _NoopClient
-from cozeloop.entities.prompt import Prompt, Message, PromptVariable
+from cozeloop.entities.prompt import Prompt, Message, PromptVariable, ExecuteResult
+from cozeloop.entities.stream import StreamReader
 from cozeloop.internal import consts, httpclient
 from cozeloop.internal.consts import ClientClosedError
 from cozeloop.internal.httpclient import Auth
@@ -269,6 +270,62 @@ def prompt_format(self, prompt: Prompt, variables: Dict[str, PromptVariable]) ->
             raise ClientClosedError()
         return self._prompt_provider.prompt_format(prompt, variables)
 
+    def execute_prompt(
+        self,
+        prompt_key: str,
+        *,
+        version: Optional[str] = None,
+        label: Optional[str] = None,
+        variable_vals: Optional[Dict[str, Any]] = None,
+        messages: Optional[List[Message]] = None,
+        stream: bool = False,
+        timeout: Optional[int] = None
+    ) -> Union[ExecuteResult, StreamReader[ExecuteResult]]:
+        """
+        执行Prompt请求
+        
+        :param timeout: 请求超时时间（秒），可选，默认为600秒（10分钟）
+        """
+        if self._closed:
+            raise ClientClosedError()
+        return self._prompt_provider.execute_prompt(
+            prompt_key,
+            version=version,
+            label=label,
+            variable_vals=variable_vals,
+            messages=messages,
+            stream=stream,
+            timeout=timeout
+        )
+
+    async def aexecute_prompt(
+        self,
+        prompt_key: str,
+        *,
+        version: Optional[str] = None,
+        label: Optional[str] = None,
+        variable_vals: Optional[Dict[str, Any]] = None,
+        messages: Optional[List[Message]] = None,
+        stream: bool = False,
+        timeout: Optional[int] = None
+    ) -> Union[ExecuteResult, StreamReader[ExecuteResult]]:
+        """
+        异步执行Prompt请求
+        
+        :param timeout: 请求超时时间（秒），可选，默认为600秒（10分钟）
+        """
+        if self._closed:
+            raise ClientClosedError()
+        return await self._prompt_provider.aexecute_prompt(
+            prompt_key,
+            version=version,
+            label=label,
+            variable_vals=variable_vals,
+            messages=messages,
+            stream=stream,
+            timeout=timeout
+        )
+
     def start_span(
             self,
             name: str,
@@ -368,6 +425,58 @@ def prompt_format(prompt: Prompt, variables: Dict[str, Any]) -> List[Message]:
     return get_default_client().prompt_format(prompt, variables)
 
 
+def execute_prompt(
+    prompt_key: str,
+    *,
+    version: Optional[str] = None,
+    label: Optional[str] = None,
+    variable_vals: Optional[Dict[str, Any]] = None,
+    messages: Optional[List[Message]] = None,
+    stream: bool = False,
+    timeout: Optional[int] = None
+) -> Union[ExecuteResult, StreamReader[ExecuteResult]]:
+    """
+    执行Prompt请求
+    
+    :param timeout: 请求超时时间（秒），可选，默认为600秒（10分钟）
+    """
+    return get_default_client().execute_prompt(
+        prompt_key,
+        version=version,
+        label=label,
+        variable_vals=variable_vals,
+        messages=messages,
+        stream=stream,
+        timeout=timeout
+    )
+
+
+async def aexecute_prompt(
+    prompt_key: str,
+    *,
+    version: Optional[str] = None,
+    label: Optional[str] = None,
+    variable_vals: Optional[Dict[str, Any]] = None,
+    messages: Optional[List[Message]] = None,
+    stream: bool = False,
+    timeout: Optional[int] = None
+) -> Union[ExecuteResult, StreamReader[ExecuteResult]]:
+    """
+    异步执行Prompt请求
+    
+    :param timeout: 请求超时时间（秒），可选，默认为600秒（10分钟）
+    """
+    return await get_default_client().aexecute_prompt(
+        prompt_key,
+        version=version,
+        label=label,
+        variable_vals=variable_vals,
+        messages=messages,
+        stream=stream,
+        timeout=timeout
+    )
+
+
 def start_span(name: str, span_type: str, *, start_time: Optional[int] = None,
                child_of: Optional[SpanContext] = None) -> Span:
     return get_default_client().start_span(name, span_type, start_time=start_time, child_of=child_of)
@@ -382,4 +491,4 @@ def get_span_from_header(header: Dict[str, str]) -> SpanContext:
 
 
 def flush() -> None:
-    return get_default_client().flush()
+    return get_default_client().flush()

cozeloop/_noop.py

@@ -2,10 +2,11 @@
 # SPDX-License-Identifier: MIT
 
 import logging
-from typing import Dict, Optional, List
+from typing import Dict, Optional, List, Union, Any
 
 from cozeloop.client import Client
-from cozeloop.entities.prompt import Prompt, Message, PromptVariable
+from cozeloop.entities.prompt import Prompt, Message, PromptVariable, ExecuteResult
+from cozeloop.entities.stream import StreamReader
 from cozeloop.internal.trace.noop_span import NoopSpan
 from cozeloop.span import SpanContext, Span
 
@@ -35,6 +36,32 @@ def prompt_format(self, prompt: Prompt, variables: Dict[str, PromptVariable]) ->
         logger.warning(f"Noop client not supported. {self.new_exception}")
         raise self.new_exception
 
+    def execute_prompt(
+        self,
+        prompt_key: str,
+        *,
+        version: Optional[str] = None,
+        label: Optional[str] = None,
+        variable_vals: Optional[Dict[str, Any]] = None,
+        messages: Optional[List[Message]] = None,
+        stream: bool = False
+    ) -> Union[ExecuteResult, StreamReader[ExecuteResult]]:
+        logger.warning(f"Noop client not supported. {self.new_exception}")
+        raise self.new_exception
+
+    async def aexecute_prompt(
+        self,
+        prompt_key: str,
+        *,
+        version: Optional[str] = None,
+        label: Optional[str] = None,
+        variable_vals: Optional[Dict[str, Any]] = None,
+        messages: Optional[List[Message]] = None,
+        stream: bool = False
+    ) -> Union[ExecuteResult, StreamReader[ExecuteResult]]:
+        logger.warning(f"Noop client not supported. {self.new_exception}")
+        raise self.new_exception
+
     def start_span(self, name: str, span_type: str, *, start_time: Optional[int] = None,
                    child_of: Optional[SpanContext] = None, start_new_trace: bool = False) -> Span:
         logger.warning(f"Noop client not supported. {self.new_exception}")
@@ -49,4 +76,4 @@ def get_span_from_header(self, header: Dict[str, str]) -> SpanContext:
         return NOOP_SPAN
 
     def flush(self) -> None:
-        logger.warning(f"Noop client not supported. {self.new_exception}")
+        logger.warning(f"Noop client not supported. {self.new_exception}")

cozeloop/entities/prompt.py

@@ -3,8 +3,8 @@
 
 from enum import Enum
 from typing import List, Optional, Union
-
-from pydantic import BaseModel, Field, ConfigDict
+from typing import List, Optional, Union, Dict, Any
+from pydantic import BaseModel
 
 
 class TemplateType(str, Enum):
@@ -47,19 +47,36 @@ class ToolChoiceType(str, Enum):
 class ContentType(str, Enum):
     TEXT = "text"
     IMAGE_URL = "image_url"
+    BASE64_DATA = "base64_data"
     MULTI_PART_VARIABLE = "multi_part_variable"
 
 
 class ContentPart(BaseModel):
     type: ContentType
     text: Optional[str] = None
     image_url: Optional[str] = None
+    base64_data: Optional[str] = None
+
+
+class FunctionCall(BaseModel):
+    name: str
+    arguments: Optional[str] = None
+
+
+class ToolCall(BaseModel):
+    index: int
+    id: str
+    type: ToolType
+    function_call: Optional[FunctionCall] = None
 
 
 class Message(BaseModel):
     role: Role
+    reasoning_content: Optional[str] = None
     content: Optional[str] = None
     parts: Optional[List[ContentPart]] = None
+    tool_call_id: Optional[str] = None
+    tool_calls: Optional[List[ToolCall]] = None
 
 
 class VariableDef(BaseModel):
@@ -109,5 +126,27 @@ class Prompt(BaseModel):
     llm_config: Optional[LLMConfig] = None
 
 
+class ExecuteParam(BaseModel):
+    """Execute参数"""
+    prompt_key: str
+    version: str = ""
+    label: str = ""
+    variable_vals: Optional[Dict[str, Any]] = None
+    messages: Optional[List[Message]] = None
+
+
+class TokenUsage(BaseModel):
+    """Token使用统计"""
+    input_tokens: int = 0
+    output_tokens: int = 0
+
+
+class ExecuteResult(BaseModel):
+    """Execute结果"""
+    message: Optional[Message] = None
+    finish_reason: Optional[str] = None
+    usage: Optional[TokenUsage] = None
+
+
 MessageLikeObject = Union[Message, List[Message]]
-PromptVariable = Union[str, MessageLikeObject]
+PromptVariable = Union[str, MessageLikeObject]

cozeloop/entities/stream.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# SPDX-License-Identifier: MIT
+
+from abc import ABC, abstractmethod
+from typing import TypeVar, Generic, AsyncIterator, Iterator
+
+T = TypeVar('T')
+
+
+class StreamReader(ABC, Generic[T]):
+    """流式读取器接口"""
+    
+    @abstractmethod
+    def __iter__(self) -> Iterator[T]:
+        """支持同步迭代 - for循环直接读取"""
+        pass
+    
+    @abstractmethod
+    def __next__(self) -> T:
+        """支持next()函数调用"""
+        pass
+    
+    @abstractmethod
+    def __aiter__(self) -> AsyncIterator[T]:
+        """支持异步迭代 - async for循环直接读取"""
+        pass
+    
+    @abstractmethod
+    async def __anext__(self) -> T:
+        """支持async next()调用"""
+        pass
+    
+    @abstractmethod
+    def close(self):
+        """关闭流"""
+        pass

cozeloop/internal/consts/__init__.py

@@ -15,6 +15,7 @@
 DEFAULT_PROMPT_CACHE_REFRESH_INTERVAL = 60
 DEFAULT_TIMEOUT = 3
 DEFAULT_UPLOAD_TIMEOUT = 30
+DEFAULT_PROMPT_EXECUTE_TIMEOUT = 600  # 10分钟，专用于execute_prompt和aexecute_prompt方法
 
 LOG_ID_HEADER = "x-tt-logid"
 AUTHORIZE_HEADER = "Authorization"
@@ -118,4 +119,4 @@
     OUTPUT: MAX_BYTES_OF_ONE_TAG_VALUE_OF_INPUT_OUTPUT,
 }
 
-BAGGAGE_SPECIAL_CHARS = {"=", ","}
+BAGGAGE_SPECIAL_CHARS = {"=", ","}