Skip to content
Snippets Groups Projects
Commit 8adc914b authored by hananel's avatar hananel
Browse files

remove hints that failed mypy

parent a1fa7419
Branches
Tags
No related merge requests found
......@@ -55,12 +55,12 @@ class Message(BaseModel):
role: str
content: str
def to_openai(self) -> dict[str, str]:
def to_openai(self):
if self.role.lower() not in ["user", "assistant", "system"]:
raise ValueError("Role must be either 'user', 'assistant' or 'system'")
return {"role": self.role, "content": self.content}
def to_cohere(self) -> dict[str, str]:
def to_cohere(self):
return {"role": self.role, "message": self.content}
......@@ -68,12 +68,12 @@ class Conversation(BaseModel):
messages: list[Message]
def split_by_topic(
self,
encoder: BaseEncoder,
threshold: float = 0.5,
split_method: Literal[
"consecutive_similarity_drop", "cumulative_similarity_drop"
] = "consecutive_similarity_drop",
self,
encoder: BaseEncoder,
threshold: float = 0.5,
split_method: Literal[
"consecutive_similarity_drop", "cumulative_similarity_drop"
] = "consecutive_similarity_drop",
) -> dict[str, list[str]]:
docs = [f"{m.role}: {m.content}" for m in self.messages]
return semantic_splitter(
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment