clean up for PR

This commit is contained in:
Brandon Hancock
2024-11-01 13:02:32 -04:00
parent 9933d8f880
commit 377793af42
2 changed files with 6 additions and 6 deletions

View File

@@ -34,6 +34,7 @@ class ContextualMemory:
formatted_results = "\n".join(
[f"- {result['context']}" for result in stm_results]
)
print("formatted_results stm", formatted_results)
return f"Recent Insights:\n{formatted_results}" if stm_results else ""
def _fetch_ltm_context(self, task) -> Optional[str]:
@@ -53,6 +54,8 @@ class ContextualMemory:
formatted_results = list(dict.fromkeys(formatted_results))
formatted_results = "\n".join([f"- {result}" for result in formatted_results]) # type: ignore # Incompatible types in assignment (expression has type "str", variable has type "list[str]")
print("formatted_results ltm", formatted_results)
return f"Historical Data:\n{formatted_results}" if ltm_results else ""
def _fetch_entity_context(self, query) -> str:
@@ -64,4 +67,5 @@ class ContextualMemory:
formatted_results = "\n".join(
[f"- {result['context']}" for result in em_results] # type: ignore # Invalid index type "str" for "str"; expected type "SupportsIndex | slice"
)
print("formatted_results em", formatted_results)
return f"Entities:\n{formatted_results}" if em_results else ""

View File

@@ -128,7 +128,7 @@ class RAGStorage(BaseRAGStorage):
) from e
class WatsonEmbeddingFunction(EmbeddingFunction):
def __call__(self, input: Documents) -> watson_models.Embeddings:
def __call__(self, input: Documents) -> Embeddings:
if isinstance(input, str):
input = [input]
@@ -147,12 +147,8 @@ class RAGStorage(BaseRAGStorage):
)
try:
print("Embedding input:", input)
embeddings = embedding.embed_documents(input)
print("Embedding output:", embeddings)
casted = cast(Embeddings, embeddings)
print("Casted:", casted)
return casted
return cast(Embeddings, embeddings)
except Exception as e:
print("Error during Watson embedding:", e)