Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -34,4 +34,5 @@ venv.bak/
/.idea
/.vscode
/output
dist/
dist/
.coda/
4 changes: 3 additions & 1 deletion cozeloop/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
close,
get_prompt,
prompt_format,
execute_prompt,
aexecute_prompt,
start_span,
get_span_from_context,
get_span_from_header,
Expand All @@ -30,4 +32,4 @@
ENV_JWT_OAUTH_PUBLIC_KEY_ID
)

from .span import SpanContext, Span
from .span import SpanContext, Span
115 changes: 112 additions & 3 deletions cozeloop/_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand All @@ -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()
33 changes: 30 additions & 3 deletions cozeloop/_noop.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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}")
Expand All @@ -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}")
45 changes: 42 additions & 3 deletions cozeloop/entities/prompt.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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]
36 changes: 36 additions & 0 deletions cozeloop/entities/stream.py
Original file line number Diff line number Diff line change
@@ -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
3 changes: 2 additions & 1 deletion cozeloop/internal/consts/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -118,4 +119,4 @@
OUTPUT: MAX_BYTES_OF_ONE_TAG_VALUE_OF_INPUT_OUTPUT,
}

BAGGAGE_SPECIAL_CHARS = {"=", ","}
BAGGAGE_SPECIAL_CHARS = {"=", ","}
Loading