Cleanups in schema and client

JoshuaC215 · Oct 5, 2024 · 3a8561f · 3a8561f
1 parent b9422be
commit 3a8561f
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/src/client/client.py b/src/client/client.py
@@ -48,7 +48,7 @@ async def ainvoke(
         async with httpx.AsyncClient() as client:
             response = await client.post(
                 f"{self.base_url}/invoke",
-                json=request.dict(),
+                json=request.model_dump(),
                 headers=self._headers,
                 timeout=self.timeout,
             )
@@ -78,7 +78,7 @@ def invoke(
             request.model = model
         response = httpx.post(
             f"{self.base_url}/invoke",
-            json=request.dict(),
+            json=request.model_dump(),
             headers=self._headers,
             timeout=self.timeout,
         )
@@ -142,7 +142,7 @@ def stream(
         with httpx.stream(
             "POST",
             f"{self.base_url}/stream",
-            json=request.dict(),
+            json=request.model_dump(),
             headers=self._headers,
             timeout=self.timeout,
         ) as response:
@@ -188,7 +188,7 @@ async def astream(
             async with client.stream(
                 "POST",
                 f"{self.base_url}/stream",
-                json=request.dict(),
+                json=request.model_dump(),
                 headers=self._headers,
                 timeout=self.timeout,
             ) as response:
@@ -215,7 +215,7 @@ async def acreate_feedback(
         async with httpx.AsyncClient() as client:
             response = await client.post(
                 f"{self.base_url}/feedback",
-                json=request.dict(),
+                json=request.model_dump(),
                 headers=self._headers,
                 timeout=self.timeout,
             )

diff --git a/src/schema/schema.py b/src/schema/schema.py
@@ -133,7 +133,9 @@ def from_langchain(cls, message: BaseMessage) -> "ChatMessage":
     def to_langchain(self) -> BaseMessage:
         """Convert the ChatMessage to a LangChain message."""
         if self.original:
-            return messages_from_dict([self.original])[0]
+            raw_original = messages_from_dict([self.original])[0]
+            raw_original.content = self.content
+            return raw_original
         match self.type:
             case "human":
                 return HumanMessage(content=self.content)