Refactor memory components to use updated Record schema

This commit is contained in:
ogabrielluiz 2024-06-11 13:05:32 -03:00
commit 60f446cb76
6 changed files with 9 additions and 13 deletions

View file

@ -1,8 +1,7 @@
from typing import Optional, cast from typing import Optional, cast
from langflow.base.memory.memory import BaseMemoryComponent from langflow.base.memory.memory import BaseMemoryComponent
from langflow.schema.schema import Record from langflow.schema.record import Record
class AstraDBMessageReaderComponent(BaseMemoryComponent): class AstraDBMessageReaderComponent(BaseMemoryComponent):

View file

@ -1,12 +1,9 @@
from typing import Optional from typing import Optional
from langchain_astradb import AstraDBChatMessageHistory
from langchain_core.messages import BaseMessage from langchain_core.messages import BaseMessage
from langflow.base.memory.memory import BaseMemoryComponent from langflow.base.memory.memory import BaseMemoryComponent
from langflow.schema.schema import Record from langflow.schema.record import Record
from langchain_core.messages import BaseMessage
class AstraDBMessageWriterComponent(BaseMemoryComponent): class AstraDBMessageWriterComponent(BaseMemoryComponent):

View file

@ -3,7 +3,7 @@ from typing import Optional, cast
from langchain_community.chat_message_histories import CassandraChatMessageHistory from langchain_community.chat_message_histories import CassandraChatMessageHistory
from langflow.base.memory.memory import BaseMemoryComponent from langflow.base.memory.memory import BaseMemoryComponent
from langflow.schema.schema import Record from langflow.schema.record import Record
class CassandraMessageReaderComponent(BaseMemoryComponent): class CassandraMessageReaderComponent(BaseMemoryComponent):

View file

@ -1,10 +1,10 @@
from typing import Optional from typing import Optional
from langflow.base.memory.memory import BaseMemoryComponent
from langflow.schema.schema import Record
from langchain_core.messages import BaseMessage
from langchain_community.chat_message_histories import CassandraChatMessageHistory from langchain_community.chat_message_histories import CassandraChatMessageHistory
from langchain_core.messages import BaseMessage
from langflow.base.memory.memory import BaseMemoryComponent
from langflow.schema.record import Record
class CassandraMessageWriterComponent(BaseMemoryComponent): class CassandraMessageWriterComponent(BaseMemoryComponent):

View file

@ -19,6 +19,6 @@ class PromptComponent(CustomComponent):
template: Prompt, template: Prompt,
**kwargs, **kwargs,
) -> Prompt: ) -> Prompt:
prompt = await Prompt.from_template_and_variables(template, kwargs) prompt = await Prompt.from_template_and_variables(template, kwargs) # type: ignore
self.status = prompt.format_text() self.status = prompt.format_text()
return prompt return prompt

View file

@ -91,4 +91,4 @@ class PythonCodeStructuredTool(CustomComponent):
tool = StructuredTool.from_function( tool = StructuredTool.from_function(
func=func, args_schema=_class, name=name, description=description, return_direct=return_direct func=func, args_schema=_class, name=name, description=description, return_direct=return_direct
) )
return tool return tool # type: ignore