langchain callback improve message init (#5)

* callback modify message init

* update tool span type

* modify tracespec package name

Author: JqRrt

cozeloop/integration/langchain/trace_callback.py

@@ -304,7 +304,7 @@ def _span_type_mapping(span_type: str) -> str:
     elif span_type == 'ReActSingleInputOutputParser':
         return 'parser'
     elif span_type == 'tool':
-        return 'plugin'
+        return 'tool'
     return span_type
 
 
@@ -406,7 +406,7 @@ def _convert_inputs(inputs: Any) -> Any:
             format_inputs['content'] = inputs.content
         return format_inputs
     if isinstance(inputs, BaseMessage):
-        message = Message(role=inputs.type, content=inputs.content)
+        message = Message(role=inputs.type, content=inputs.content, tool_calls=inputs.additional_kwargs.get('tool_calls', []))
         return message
     if isinstance(inputs, ChatPromptValue):
         return _convert_inputs(inputs.messages)

cozeloop/integration/langchain/trace_model/llm_model.py

@@ -3,7 +3,7 @@
 
 import json
 import time
-from typing import List, Optional, Union
+from typing import List, Optional, Union, Dict, Any
 from pydantic.dataclasses import dataclass
 from langchain_core.messages import BaseMessage, ToolMessage, AIMessageChunk
 from langchain_core.outputs import Generation, ChatGeneration
@@ -14,7 +14,7 @@ class ToolFunction:
     name: Optional[str] = None
     description: Optional[str] = None
     parameters: Optional[dict] = None
-    arguments: Optional[dict] = None
+    arguments: Optional[Union[dict, str]] = None
 
 
 @dataclass
@@ -50,20 +50,31 @@ class Message:
     tool_calls: List[ToolCall] = None
 
     def __post_init__(self):
-        if self.role is not None and self.role == 'AIMessageChunk':
+        if self.role is not None and (self.role == 'AIMessageChunk' or self.role == 'ai'):
             self.role = 'assistant'
         parts: Optional[List[Parts]] = []
         if isinstance(self.content, List) and all(isinstance(x, dict) for x in self.content):
+            is_parts = False
             for each in self.content:
                 text = each.get('text', None)
                 url = each.get('url', each.get('image_url', {}).get('url', None))
+                if text is None and url is None:
+                    continue
+                is_parts = True
                 parts.append(Parts(type=each.get('type', ''), text=text, image_url=ImageUrl(url=url) if url is not None else None))
-            self.content = None
+            if is_parts:
+                self.content = None
+            else:
+                self.content = self.content.__str__()
         elif isinstance(self.content, dict):
+            is_part = False
             text = self.content.get('text', None)
             url = self.content.get('url', self.content.get('image_url', {}).get('url', None))
-            parts.append(Parts(type=self.content.get('type', ''), text=text, image_url=ImageUrl(url=url) if url is not None else None))
-            self.content = None
+            if text is not None or url is not None:
+                parts.append(Parts(type=self.content.get('type', ''), text=text, image_url=ImageUrl(url=url) if url is not None else None))
+                self.content = None
+            else:
+                self.content = self.content.__str__()
         elif isinstance(self.content, List) and all(type(x, Parts) for x in self.content):
             parts = self.content
             self.content = None
@@ -172,4 +183,4 @@ def convert_tool_calls(tool_calls: list) -> List[ToolCall]:
     for tool_call in tool_calls:
         function = ToolFunction(name=tool_call.get('function', {}).get('name', ''), arguments=json.loads(tool_call.get('function', {}).get('arguments', {})))
         format_tool_calls.append(ToolCall(id=tool_call.get('id', ''), type=tool_call.get('type', ''), function=function))
-    return format_tool_calls
+    return format_tool_calls

cozeloop/integration/langchain/trace_model/runtime.py

@@ -5,12 +5,12 @@
 import importlib.metadata as metadata
 from typing import Optional, Any
 
-from cozeloop.spec import tracespce
+from cozeloop.spec import tracespec
 
 
-class RuntimeInfo(tracespce.Runtime):
-    language: Optional[str] = tracespce.V_LANG_PYTHON
-    library: Optional[str] = tracespce.V_LIB_LANGCHAIN
+class RuntimeInfo(tracespec.Runtime):
+    language: Optional[str] = tracespec.V_LANG_PYTHON
+    library: Optional[str] = tracespec.V_LIB_LANGCHAIN
 
     def model_post_init(self, context: Any) -> None:
         try:

cozeloop/internal/consts/__init__.py

@@ -4,7 +4,7 @@
 from contextvars import ContextVar
 
 from .error import *
-from ...spec.tracespce.span_key import INPUT_TOKENS, OUTPUT_TOKENS, TOKENS, INPUT, OUTPUT
+from ...spec.tracespec.span_key import INPUT_TOKENS, OUTPUT_TOKENS, TOKENS, INPUT, OUTPUT
 
 # default values for loop client
 # ComBaseURL = "https://api.coze.com"

cozeloop/internal/prompt/converter.py

@@ -3,7 +3,7 @@
 
 from typing import List, Dict
 
-from cozeloop.spec.tracespce import PromptInput, PromptOutput, ModelMessage, PromptArgument
+from cozeloop.spec.tracespec import PromptInput, PromptOutput, ModelMessage, PromptArgument
 from cozeloop.entities.prompt import (
     Prompt as EntityPrompt,
     Message as EntityMessage,

cozeloop/internal/prompt/prompt.py

@@ -7,7 +7,7 @@
 from jinja2 import Environment, BaseLoader, Undefined
 from jinja2.utils import missing, object_type_repr
 
-from cozeloop.spec.tracespce import PROMPT_KEY, INPUT, PROMPT_VERSION, V_SCENE_PROMPT_TEMPLATE, V_SCENE_PROMPT_HUB
+from cozeloop.spec.tracespec import PROMPT_KEY, INPUT, PROMPT_VERSION, V_SCENE_PROMPT_TEMPLATE, V_SCENE_PROMPT_HUB
 from cozeloop.entities.prompt import (Prompt, Message, VariableDef, VariableType, TemplateType, Role,
                                       PromptVariable)
 from cozeloop.internal import consts

cozeloop/internal/trace/exporter.py

@@ -6,7 +6,7 @@
 import time
 from typing import Dict, List, Optional, Tuple, Callable, Any
 
-from cozeloop.spec.tracespce import ModelInput, ModelMessagePart, ModelMessagePartType, ModelImageURL, ModelFileURL, ModelOutput
+from cozeloop.spec.tracespec import ModelInput, ModelMessagePart, ModelMessagePartType, ModelImageURL, ModelFileURL, ModelOutput
 from cozeloop.internal.consts import *
 from cozeloop.internal.httpclient import Client, BaseResponse
 from cozeloop.internal.trace.model.model import UploadType, Attachment, ObjectStorage

cozeloop/internal/trace/noop_span.py

@@ -8,7 +8,7 @@
 
 from cozeloop.entities.prompt import Prompt
 from cozeloop.span import Span
-from cozeloop.spec.tracespce import Runtime
+from cozeloop.spec.tracespec import Runtime
 
 
 class NoopSpan(Span, ABC):

cozeloop/internal/trace/span.py

@@ -10,7 +10,7 @@
 import urllib.parse
 
 from cozeloop import span
-from cozeloop.spec.tracespce import (ModelInput, ModelOutput, ModelMessagePartType, ModelMessage, ModelMessagePart,
+from cozeloop.spec.tracespec import (ModelInput, ModelOutput, ModelMessagePartType, ModelMessage, ModelMessagePart,
                                      ModelImageURL, ModelFileURL, ModelChoice, Runtime, ERROR, PROMPT_KEY,
                                      PROMPT_VERSION, MODEL_PROVIDER, MODEL_NAME, RUNTIME_, CALL_OPTIONS,
                                      V_SCENE_CUSTOM, V_LANG_PYTHON)

cozeloop/internal/trace/trace.py

@@ -5,7 +5,7 @@
 from datetime import datetime
 from typing import Dict, Optional
 
-from cozeloop.spec.tracespce import Runtime, RUNTIME_
+from cozeloop.spec.tracespec import Runtime, RUNTIME_
 from cozeloop.internal import consts
 from cozeloop.internal.trace.span import from_header, Span, SpanContext, \
     get_newest_span_from_context, set_span_to_context, logger