Merge pull request #61 from log10-io/wenzhe/convert_message

refactor func to convert log10 messages

examples/logging/langchain_model_logger.py

@@ -1,7 +1,7 @@
 from langchain import OpenAI
 from langchain.chat_models import ChatAnthropic
 from langchain.chat_models import ChatOpenAI
-from langchain.schema import HumanMessage
+from langchain.schema import AIMessage, HumanMessage, SystemMessage
 
 from log10.langchain import Log10Callback
 from log10.llm import Log10Config
@@ -11,8 +11,10 @@
 
 
 messages = [
-    HumanMessage(content="You are a ping pong machine"),
+    SystemMessage(content="You are a ping pong machine"),
     HumanMessage(content="Ping?"),
+    AIMessage(content="Pong"),
+    HumanMessage(content="Ping ping"),
 ]
 
 llm = ChatOpenAI(

log10/langchain.py

@@ -37,6 +37,21 @@ def kwargs_to_hparams(kwargs: Dict[str, Any]) -> Dict[str, Any]:
     return hparams
 
 
+def get_log10_messages(langchain_messages: List[BaseMessage]) -> List[Message]:
+    role_map = {AIMessage: "assistant", HumanMessage: "user", SystemMessage: "system"}
+
+    for m in langchain_messages:
+        logger.debug(f"message: {m}")
+        if type(m) not in role_map:
+            raise BaseException(
+                f"Unsupported message type {type(m)}. Supported types: {list(role_map.values())}"
+            )
+
+    return [
+        Message(role=role_map[type(m)], content=m.content) for m in langchain_messages
+    ]
+
+
 class Log10Callback(BaseCallbackHandler, LLM):
     """Callback Handler that prints to std out."""
 
@@ -122,19 +137,7 @@ def on_chat_model_start(
             raise BaseException("Only support one message at a time")
 
         # Convert messages to log10 format
-        log10_messages = []
-        for message in messages[0]:
-            logger.debug(f"message: {message}")
-            if isinstance(message, HumanMessage):
-                log10_messages.append(Message(role="user", content=message.content))
-            elif isinstance(message, AIMessage):
-                log10_messages.append(
-                    Message(role="assistant", content=message.content)
-                )
-            elif isinstance(message, SystemMessage):
-                log10_messages.append(Message(role="system", content=message.content))
-            else:
-                raise BaseException(f"Unknown message type {type(message)}")
+        log10_messages = get_log10_messages(messages[0])
 
         request = {
             "messages": [message.to_dict() for message in log10_messages],

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api"
 [tool.poetry]
 name = "log10-io"
 
-version = "0.4.1"
+version = "0.4.3"
 authors = ["log10 team"]
 license = "MIT"
 description = "Unified LLM data management"

setup.py

@@ -5,7 +5,7 @@
 
 setup(
     name="Log10",
-    version="0.4.1",
+    version="0.4.3",
     description="Log10 LLM data management",
     author="Log10 team",
     author_email="team@log10.io",