From 4df1fc492223cc8a0360e00eaf0e6efac2c20b3e Mon Sep 17 00:00:00 2001 From: Benjoyo Date: Tue, 2 Jan 2024 16:07:15 +0100 Subject: [PATCH] fix test warnings --- bpm-ai-core/.gitignore | 1 + bpm-ai-core/bpm_ai_core/llm/common/llm.py | 8 +++----- bpm-ai-core/bpm_ai_core/llm/common/message.py | 10 ++++------ bpm-ai-core/bpm_ai_core/tracing/config.py | 3 +++ bpm-ai-core/bpm_ai_core/tracing/tracing.py | 14 ++++++-------- bpm-ai-core/tests/pytest.ini | 2 +- 6 files changed, 18 insertions(+), 20 deletions(-) create mode 100644 bpm-ai-core/bpm_ai_core/tracing/config.py diff --git a/bpm-ai-core/.gitignore b/bpm-ai-core/.gitignore index e796581..7399b5b 100644 --- a/bpm-ai-core/.gitignore +++ b/bpm-ai-core/.gitignore @@ -32,6 +32,7 @@ MANIFEST.in .venv .env +.env.test /releases/* pip-wheel-metadata /poetry.toml diff --git a/bpm-ai-core/bpm_ai_core/llm/common/llm.py b/bpm-ai-core/bpm_ai_core/llm/common/llm.py index b68fe7b..1dfd285 100644 --- a/bpm-ai-core/bpm_ai_core/llm/common/llm.py +++ b/bpm-ai-core/bpm_ai_core/llm/common/llm.py @@ -6,7 +6,7 @@ from bpm_ai_core.llm.common.message import ChatMessage from bpm_ai_core.llm.common.tool import Tool from bpm_ai_core.prompt.prompt import Prompt -from bpm_ai_core.tracing.tracing import LangsmithTracer +from bpm_ai_core.tracing.config import tracer class LLM(ABC): @@ -19,8 +19,6 @@ class LLM(ABC): max_retries: int = 0 retryable_exceptions: List[Type[BaseException]] = [Exception] - tracer = LangsmithTracer() - @retry( wait=wait_exponential(multiplier=1.5, min=2, max=60), stop=stop_after_attempt(max_retries), @@ -37,9 +35,9 @@ def predict( messages = prompt.format(llm_name=self.name()) - self.tracer.start_llm_trace(self, messages, self.predict.retry.statistics['attempt_number'], tools) + tracer.start_llm_trace(self, messages, self.predict.retry.statistics['attempt_number'], tools) completion = self._predict(messages, output_schema, tools) - self.tracer.end_llm_trace(completion) + tracer.end_llm_trace(completion) return completion diff --git a/bpm-ai-core/bpm_ai_core/llm/common/message.py b/bpm-ai-core/bpm_ai_core/llm/common/message.py index f4a401d..b5137bb 100644 --- a/bpm-ai-core/bpm_ai_core/llm/common/message.py +++ b/bpm-ai-core/bpm_ai_core/llm/common/message.py @@ -4,13 +4,14 @@ from typing import Optional, Literal, Any, Union, List from PIL import Image -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, ConfigDict from bpm_ai_core.llm.common.tool import Tool +from bpm_ai_core.tracing.config import tracer class ChatMessage(BaseModel): - content: Optional[Union[str, dict, List[Union[str, Image]]]] = None + content: Optional[Union[str, dict, List[Union[str, Image.Image]]]] = None """ The contents of the message. Either a string for normal completions, @@ -29,8 +30,7 @@ class ChatMessage(BaseModel): The name of the author of this message. """ - class Config: - arbitrary_types_allowed = True + model_config = ConfigDict(arbitrary_types_allowed=True) class SingleToolCallMessage(BaseModel): @@ -65,8 +65,6 @@ def invoke(self): async def ainvoke(self) -> Any: _callable = self.tool.callable inputs = self.payload_dict() - from bpm_ai_core.tracing.tracing import LangsmithTracer - tracer = LangsmithTracer() # todo tracer.start_function_trace(self.tool, inputs) if inspect.iscoroutinefunction(_callable): result = await _callable(**inputs) diff --git a/bpm-ai-core/bpm_ai_core/tracing/config.py b/bpm-ai-core/bpm_ai_core/tracing/config.py new file mode 100644 index 0000000..d2c4f21 --- /dev/null +++ b/bpm-ai-core/bpm_ai_core/tracing/config.py @@ -0,0 +1,3 @@ +from bpm_ai_core.tracing.tracing import LangsmithTracer + +tracer = LangsmithTracer() \ No newline at end of file diff --git a/bpm-ai-core/bpm_ai_core/tracing/tracing.py b/bpm-ai-core/bpm_ai_core/tracing/tracing.py index 7945974..ac82164 100644 --- a/bpm-ai-core/bpm_ai_core/tracing/tracing.py +++ b/bpm-ai-core/bpm_ai_core/tracing/tracing.py @@ -13,10 +13,6 @@ from langsmith.run_helpers import LangSmithExtra, _collect_extra, _TraceableContainer, _PROJECT_NAME, \ _PARENT_RUN_TREE, _TAGS, _METADATA -from bpm_ai_core.llm.common.message import ChatMessage, ToolCallsMessage -from bpm_ai_core.llm.common.tool import Tool -from bpm_ai_core.util.openai import messages_to_openai_dicts, json_schema_to_openai_function - class LangsmithTracer: @@ -27,10 +23,11 @@ def __init__(self): def start_llm_trace( self, llm: Any, - messages: List[ChatMessage], + messages: List, current_try: int, - tools: Optional[List[Tool]] = None + tools: Optional[List] = None ): + from bpm_ai_core.util.openai import messages_to_openai_dicts, json_schema_to_openai_function inputs = { "messages": messages_to_openai_dicts(messages), "model": llm.model, @@ -45,7 +42,8 @@ def start_llm_trace( inputs=inputs ) - def end_llm_trace(self, completion: Optional[ChatMessage] = None, error_msg: Optional[str] = None): + def end_llm_trace(self, completion = None, error_msg: Optional[str] = None): + from bpm_ai_core.llm.common.message import ToolCallsMessage choices = { "choices": [{ "message": { @@ -64,7 +62,7 @@ def end_llm_trace(self, completion: Optional[ChatMessage] = None, error_msg: Opt def start_function_trace( self, - function: Tool, + function, inputs: dict, ): self.start_trace( diff --git a/bpm-ai-core/tests/pytest.ini b/bpm-ai-core/tests/pytest.ini index 820fc5e..09ddc1d 100644 --- a/bpm-ai-core/tests/pytest.ini +++ b/bpm-ai-core/tests/pytest.ini @@ -1,2 +1,2 @@ [pytest] -env_files = .env +env_files = .env.test