Add from_lc_message method to Record class in schema.py

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-04-16 13:22:31 -03:00
commit dd7e5a7ab5

View file

@ -2,6 +2,7 @@ import copy
from typing import Literal, Optional
from langchain_core.documents import Document
from langchain_core.messages import BaseMessage
from pydantic import BaseModel, model_validator
@ -54,6 +55,21 @@ class Record(BaseModel):
data["text"] = document.page_content
return cls(data=data, text_key="text")
@classmethod
def from_lc_message(cls, message: BaseMessage) -> "Record":
"""
Converts a BaseMessage to a Record.
Args:
message (BaseMessage): The BaseMessage to convert.
Returns:
Record: The converted Record.
"""
data = {"text": message.content}
data["metadata"] = message.to_json()
return cls(data=data, text_key="text")
def __add__(self, other: "Record") -> "Record":
"""
Combines the data of two records by attempting to add values for overlapping keys