Fix langchain runtime, change examples.

Change-Id: Ie30079d23e168c3b79faf8fae5beaec8d1128ad1

Author: taoyifan89

cozeloop/integration/langchain/trace_callback.py

@@ -3,7 +3,6 @@
 
 from __future__ import annotations
 import json
-import threading
 import time
 import traceback
 from typing import List, Dict, Union, Any, Optional
@@ -23,7 +22,8 @@
 from cozeloop.integration.langchain.trace_model.runtime import RuntimeInfo
 from cozeloop.integration.langchain.util import calc_token_usage, get_prompt_tag
 
-_trace_callback_client = None
+_trace_callback_client: Optional[Client] = None
+
 
 class LoopTracer:
     @classmethod
@@ -48,16 +48,15 @@ def __init__(self):
 
     def on_llm_start(self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any) -> Any:
         span_tags = {}
+        span_name = serialized.get('name', 'unknown')
+
+        flow_span = self._new_flow_span(span_name, 'model', **kwargs)
         try:
             span_tags['input'] = ModelTraceInput([BaseMessage(type='', content=prompt) for prompt in prompts],
                                                  kwargs.get('invocation_params', {})).to_json()
-            span_name = serialized['name']
         except Exception as e:
-            span_name = 'unknown'
-            span_tags['internal_error'] = repr(e)
-            span_tags['internal_error_trace'] = traceback.format_exc()
+            flow_span.set_error(e)
         finally:
-            flow_span = self._new_flow_span(span_name, 'model', **kwargs)
             span_tags.update(_get_model_span_tags(**kwargs))
             self._set_span_tags(flow_span, span_tags)
             #  Store some pre-aspect information.
@@ -67,15 +66,14 @@ def on_llm_start(self, serialized: Dict[str, Any], prompts: List[str], **kwargs:
 
     def on_chat_model_start(self, serialized: Dict[str, Any], messages: List[List[BaseMessage]], **kwargs: Any) -> Any:
         span_tags = {}
+        span_name = serialized.get('name', 'unknown')
+
+        flow_span = self._new_flow_span(span_name, 'model', **kwargs)
         try:
             span_tags['input'] = ModelTraceInput(messages, kwargs.get('invocation_params', {})).to_json()
-            span_name = serialized['name']
         except Exception as e:
-            span_name = 'unknown'
-            span_tags['internal_error'] = repr(e)
-            span_tags['internal_error_trace'] = traceback.format_exc()
+            flow_span.set_error(e)
         finally:
-            flow_span = self._new_flow_span(span_name, 'model', **kwargs)
             span_tags.update(_get_model_span_tags(**kwargs))
             self._set_span_tags(flow_span, span_tags)
             #  Store some pre-aspect information.
@@ -105,13 +103,8 @@ def on_llm_end(self, response: LLMResult, **kwargs: Any) -> Any:
             if run_info is not None and run_info.model_meta is not None:
                 model_name = run_info.model_meta.model_name
                 input_messages = run_info.model_meta.message
-                token_usage = {
-                    'input_tokens': calc_token_usage(input_messages, model_name),
-                    'output_tokens': calc_token_usage(response, model_name),
-                    'tokens': 0
-                }
-                token_usage['tokens'] = token_usage['input_tokens'] + token_usage['output_tokens']
-                self._set_span_tags(flow_span, token_usage, need_convert_tag_value=False)
+                flow_span.set_input_tokens(calc_token_usage(input_messages, model_name))
+                flow_span.set_output_tokens(calc_token_usage(response, model_name))
         # finish flow_span
         flow_span.finish()
 
@@ -139,24 +132,17 @@ def on_chain_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: A
         if flow_span is None:
             span_name = '_Exception' if isinstance(error, Exception) else '_KeyboardInterrupt'
             flow_span = self._new_flow_span(span_name, 'chain_error', **kwargs)
-        flow_span.set_tags({'error': repr(error)})
+        flow_span.set_error(error)
         flow_span.set_tags({'error_trace': traceback.format_exc()})
-        flow_span.set_tags({'_status_code': -1})
         flow_span.finish()
 
     def on_tool_start(
             self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
     ) -> Any:
         span_tags = {'input': input_str, **serialized}
-        try:
-            span_name = serialized['name']
-        except Exception as e:
-            span_name = 'unknown'
-            span_tags['internal_error'] = repr(e)
-            span_tags['internal_error_trace'] = traceback.format_exc()
-        finally:
-            flow_span = self._new_flow_span(span_name, 'tool', **kwargs)
-            self._set_span_tags(flow_span, span_tags)
+        span_name = serialized.get('name', 'unknown')
+        flow_span = self._new_flow_span(span_name, 'tool', **kwargs)
+        self._set_span_tags(flow_span, span_tags)
 
     def on_tool_end(self, output: str, **kwargs: Any) -> Any:
         flow_span = self._get_flow_span(**kwargs)
@@ -170,9 +156,8 @@ def on_tool_error(
         if flow_span is None:
             span_name = '_Exception' if isinstance(error, Exception) else '_KeyboardInterrupt'
             flow_span = self._new_flow_span(span_name, 'tool_error', **kwargs)
-        flow_span.set_tags({'error': repr(error)})
+        flow_span.set_error(error)
         flow_span.set_tags({'error_trace': traceback.format_exc()})
-        flow_span.set_tags({'_status_code': -1})
         flow_span.finish()
 
     def on_text(self, text: str, **kwargs: Any) -> Any:
@@ -241,9 +226,7 @@ def _new_flow_span(self, span_name: str, span_type: str, **kwargs: Any) -> Span:
         run_id = str(kwargs['run_id'])
         self.run_map[run_id] = Run(run_id, flow_span, span_type)
         # set default tags
-        # flow_span.set_tags({'space_id': self._space_id})
-        flow_span.set_tags({'span_type': span_type})
-        flow_span.set_tags({'runtime': RuntimeInfo().to_json()})
+        flow_span.set_runtime(RuntimeInfo())
         return flow_span
 
     def _get_flow_span(self, **kwargs: Any) -> Span:
@@ -252,13 +235,6 @@ def _get_flow_span(self, **kwargs: Any) -> Span:
             return self.run_map[run_id].span
         return None
 
-    def _set_internal_error_span(self, error: Exception, **kwargs: Any) -> None:
-        flow_span = self._new_flow_span('internal_error', 'error', **kwargs)
-        flow_span.set_tags({'internal_error': error})
-        flow_span.set_tags({'internal_error_trace': traceback.format_exc()})
-        flow_span.set_tags({'_status_code': -1})
-        flow_span.finish()
-
     def _set_span_tags(self, flow_span: Span, tags: Dict[str, Any], need_convert_tag_value=True) -> None:
         for key, value in tags.items():
             report_value = value

cozeloop/integration/langchain/trace_model/runtime.py

@@ -2,36 +2,23 @@
 # SPDX-License-Identifier: MIT
 
 import json
-import platform as platform_pkg
 import importlib.metadata as metadata
-from typing import Optional
-from pydantic.dataclasses import dataclass
+from typing import Optional, Any
 
-from cozeloop.internal.version import VERSION
+from cozeloop.spec import tracespce
 
-@dataclass
-class RuntimeInfo:
-    language: Optional[str] = 'python'
-    library: Optional[str] = 'langchain'
-    runtime: Optional[str] = 'python'
-    runtime_version: Optional[str] = platform_pkg.python_version()
-    py_implementation: Optional[str] = platform_pkg.python_implementation()
-    loop_sdk_version: Optional[str] = None
-    langchain_version: Optional[str] = None
-    langchain_core_version: Optional[str] = None
 
-    def __post_init__(self):
+class RuntimeInfo(tracespce.Runtime):
+    language: Optional[str] = tracespce.V_LANG_PYTHON
+    library: Optional[str] = tracespce.V_LIB_LANGCHAIN
+
+    def model_post_init(self, context: Any) -> None:
         try:
             langchain_version = metadata.version('langchain')
         except metadata.PackageNotFoundError:
             langchain_version = ''
-        try:
-            langchain_core_version = metadata.version('langchain-core')
-        except metadata.PackageNotFoundError:
-            langchain_core_version = ''
-        self.loop_sdk_version = VERSION
-        self.langchain_version = langchain_version
-        self.langchain_core_version = langchain_core_version
+
+        self.library_version = langchain_version
 
     def to_json(self):
         return json.dumps(

cozeloop/internal/consts/__init__.py

@@ -61,8 +61,8 @@
 COMMA = ","
 
 # On the basis of W3C, the "loop" prefix is added to avoid conflicts with other traces that use W3C.
-TRACE_CONTEXT_HEADER_PARENT = "X-Loop-Traceparent"
-TRACE_CONTEXT_HEADER_BAGGAGE = "X-Loop-Tracestate"
+TRACE_CONTEXT_HEADER_PARENT = "X-Cozeloop-Traceparent"
+TRACE_CONTEXT_HEADER_BAGGAGE = "X-Cozeloop-Tracestate"
 
 TRACE_PROMPT_HUB_SPAN_TYPE = "prompt_hub"
 TRACE_PROMPT_TEMPLATE_SPAN_TYPE = "prompt"

cozeloop/internal/httpclient/user_agent.py

@@ -9,12 +9,12 @@
 from cozeloop import internal
 
 # User agent components
-USER_AGENT_SDK = "loop-python"
+USER_AGENT_SDK = "cozeloop-python"
 USER_AGENT_LANG = "python"
 USER_AGENT_LANG_VERSION = platform.python_version()
 USER_AGENT_OS_NAME = platform.system().lower()
 USER_AGENT_OS_VERSION = os.getenv("OSVERSION", "unknown")
-SCENE = "loop"
+SCENE = "cozeloop"
 SOURCE = "openapi"
 
 USER_AGENT = f"{USER_AGENT_SDK}/{internal.__version__} {USER_AGENT_LANG}/{USER_AGENT_LANG_VERSION} {USER_AGENT_OS_NAME}/{USER_AGENT_OS_VERSION}"

cozeloop/internal/trace/noop_span.py

@@ -8,6 +8,7 @@
 
 from cozeloop.entities.prompt import Prompt
 from cozeloop.span import Span
+from cozeloop.spec.tracespce import Runtime
 
 
 class NoopSpan(Span, ABC):
@@ -106,6 +107,9 @@ def set_output_tokens(self, output_tokens: int) -> None:
     def set_start_time_first_resp(self, start_time_first_resp: int) -> None:
         pass
 
+    def set_runtime(self, runtime: Runtime) -> None:
+        pass
+
     def __enter__(self):
         return self
 

cozeloop/internal/trace/span.py

@@ -342,6 +342,13 @@ def set_output_tokens(self, output_tokens: int):
     def set_start_time_first_resp(self, start_time_first_resp: int):
         self.set_tags({START_TIME_FIRST_RESP: start_time_first_resp})
 
+    def set_runtime(self, runtime: Runtime) -> None:
+        r = Runtime(scene=V_SCENE_CUSTOM, library=runtime.library, library_version=runtime.library_version)
+        with self.lock:
+            if self.system_tag_map is None:
+                self.system_tag_map = {}
+            self.system_tag_map[RUNTIME_] = r
+
     def get_rectified_map(self, input_map: Dict[str, Any]) -> (Dict[str, Any], List[str], int):
         validate_map = {}
         cut_off_keys = []

cozeloop/span.py

@@ -6,6 +6,7 @@
 from typing import Dict, Any
 from datetime import datetime
 from cozeloop.entities.prompt import Prompt
+from cozeloop.spec.tracespce import Runtime
 
 
 class SpanContext(ABC):
@@ -159,6 +160,12 @@ def set_start_time_first_resp(self, start_time_first_resp: int) -> None:
         Key: `start_time_first_resp`.
         """
 
+    @abstractmethod
+    def set_runtime(self, runtime: Runtime) -> None:
+        """
+        Set the runtime of the span. Only used for integration.
+        Key: `runtime`.
+        """
 
 class Span(CommonSpanSetter, SpanContext):
     """

cozeloop/spec/tracespce/span_value.py

@@ -33,6 +33,7 @@
 V_SCENE_CUSTOM = "custom"          # user custom, it has the same meaning as blank.
 V_SCENE_PROMPT_HUB = "prompt_hub"      # get_prompt
 V_SCENE_PROMPT_TEMPLATE = "prompt_template"  # prompt_template
+V_SCENE_INTEGRATION = "integration"   # integration like langchain
 
 # Tag values for prompt input.
 V_PROMPT_ARG_SOURCE_INPUT = "input"

examples/lcel/lcel.py

@@ -2,12 +2,9 @@
 # SPDX-License-Identifier: MIT
 
 import logging
-import os
-import time
 
-from langchain.callbacks.tracers import ConsoleCallbackHandler
 from langchain_core.runnables import RunnableConfig
-from langchain_openai import AzureChatOpenAI
+from langchain_openai import ChatOpenAI
 from langchain_core.output_parsers import StrOutputParser
 
 from cozeloop import set_log_level, new_client
@@ -18,30 +15,27 @@
 def do_lcel_demo():
     # Configure the parameters for the large model. The keys in os.environ are standard keys for Langchain and must be
     # followed. This is just a demo, and the connectivity of the large model needs to be ensured by the user.
-    # os.environ['AZURE_OPENAI_API_KEY'] = 'xxx'  # need set a llm api key
-    # os.environ['OPENAI_API_VERSION'] = '2024-05-13'  # llm version, see more: https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#rest-api-versioning
-    # os.environ['AZURE_OPENAI_ENDPOINT'] = 'https://xxx'  # llm endpoint
-    # os.environ['AUZURE_DEPLOYMENT'] = 'gpt-4o-2024-05-13'
+    # os.environ['OPENAI_API_KEY'] = 'xxx'  # need set a openai key
 
-    # Configure the Loop environment variables. This is just a demo, and the keys in os.environ are not for reference.
+    # Configure the CozeLoop environment variables. This is just a demo, and the keys in os.environ are not for reference.
     # The specific implementation method is determined by the business side.
     # Set the following environment variables first (Assuming you are using a PAT token.).
     # os.environ['COZELOOP_API_TOKEN'] = 'your token'
     # os.environ['COZELOOP_WORKSPACE_ID'] = 'your workspace id'
 
-    client = new_client(ultra_large_report=True)
+    client = new_client()
     trace_callback_handler = LoopTracer.get_callback_handler(client)
     # init llm model
-    llm_model = AzureChatOpenAI(azure_deployment=os.environ['AUZURE_DEPLOYMENT'])
+    llm_model = ChatOpenAI(model="doubao-1-5-vision-pro-32k-250115", base_url="https://ark.cn-beijing.volces.com/api/v3")
 
     # execute lcel, and print intermediate results.
     lcel_sequence = llm_model | StrOutputParser()
     output = lcel_sequence.invoke(
         input='用你所学的技巧，帮我生成几个有意思的问题',
-        config=RunnableConfig(callbacks=[ConsoleCallbackHandler(), trace_callback_handler])
+        config=RunnableConfig(callbacks=[trace_callback_handler])
     )
-    time.sleep(5) # async report, so sleep wait for report finish
     print('\n====== model output start ======\n' + output + '\n====== model output finish ======\n')
+    client.close()
 
 
 if __name__ == "__main__":

examples/lcel/lcel_stream.py

@@ -5,36 +5,32 @@
 import os
 import time
 
-from langchain.callbacks.tracers import ConsoleCallbackHandler
 from langchain_core.runnables import RunnableConfig
-from langchain_openai import AzureChatOpenAI
+from langchain_openai import ChatOpenAI
 from langchain_core.output_parsers import StrOutputParser
 
-from cozeloop import set_log_level
+from cozeloop import set_log_level, new_client
 from cozeloop.integration.langchain.trace_callback import LoopTracer
 
 logger = logging.getLogger(__name__)
 
 def do_lcel_stream_demo():
     # Configure the parameters for the llm. The keys in os.environ are standard keys for Langchain and must be
     # followed. This is just a demo, and the connectivity of the llm needs to be ensured by the user.
-    # os.environ['AZURE_OPENAI_API_KEY'] = 'xxx'  # need set a llm api key
-    # os.environ['OPENAI_API_VERSION'] = '2024-05-13'  # llm version, see more: https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#rest-api-versioning
-    # os.environ['AZURE_OPENAI_ENDPOINT'] = 'https://xxx'  # llm endpoint
-    # os.environ['AUZURE_DEPLOYMENT'] = 'gpt-4o-2024-05-13'
+    # os.environ['OPENAI_API_KEY'] = 'xxx'  # need set a openai key
 
-    # Configure the Loop environment variables. This is just a demo, and the keys in os.environ are not for reference.
+    # Configure the CozeLoop environment variables. This is just a demo, and the keys in os.environ are not for reference.
     # The specific implementation method is determined by the business side.
-
     # Set the following environment variables first (Assuming you are using a PAT token.).
     # COZELOOP_WORKSPACE_ID=your workspace id
     # COZELOOP_API_TOKEN=your token
     # os.environ['COZELOOP_API_TOKEN'] = 'your token'
     # os.environ['COZELOOP_WORKSPACE_ID'] = 'your workspace'
 
-    trace_callback_handler = LoopTracer.get_callback_handler()
+    client = new_client()
+    trace_callback_handler = LoopTracer.get_callback_handler(client)
     # init llm model
-    llm_model = AzureChatOpenAI(azure_deployment=os.environ['AUZURE_DEPLOYMENT'])
+    llm_model = ChatOpenAI(model="doubao-1-5-vision-pro-32k-250115", base_url="https://ark.cn-beijing.volces.com/api/v3")
 
     # execute lcel, and print intermediate results.
     lcel_sequence = llm_model | StrOutputParser()
@@ -46,8 +42,8 @@ def do_lcel_stream_demo():
         chunks.append(chunk)
         print(chunk, end='', flush=True)
 
-    time.sleep(5) # async report, so sleep wait for report finish
     print('\n====== model output start ======\n' + ''.join(chunks) + '\n====== model output finish ======\n')
+    client.close()
 
 
 if __name__ == "__main__":