fix: collection [bull_memory] already exists

This commit is contained in:
Phuc Nguyen Minh 2025-06-14 09:19:48 +07:00 committed by GitHub
parent 570644d939
commit 39ef79f39a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 6 deletions

View File

@ -1,20 +1,17 @@
import chromadb import chromadb
from chromadb.config import Settings from chromadb.config import Settings
from openai import OpenAI from openai import OpenAI
import numpy as np
class FinancialSituationMemory: class FinancialSituationMemory:
def __init__(self, name): def __init__(self, name):
self.client = OpenAI() self.client = OpenAI()
self.chroma_client = chromadb.Client(Settings(allow_reset=True)) self.chroma_client = chromadb.Client(Settings(allow_reset=True))
self.situation_collection = self.chroma_client.create_collection(name=name) self.situation_collection = self.chroma_client.get_or_create_collection(name=name)
def get_embedding(self, text): def get_embedding(self, text):
"""Get OpenAI embedding for a text""" """Get OpenAI embedding for a text"""
response = self.client.embeddings.create( response = self.client.embeddings.create(model="text-embedding-ada-002", input=text)
model="text-embedding-ada-002", input=text
)
return response.data[0].embedding return response.data[0].embedding
def add_situations(self, situations_and_advice): def add_situations(self, situations_and_advice):
@ -92,7 +89,7 @@ if __name__ == "__main__":
# Example query # Example query
current_situation = """ current_situation = """
Market showing increased volatility in tech sector, with institutional investors Market showing increased volatility in tech sector, with institutional investors
reducing positions and rising interest rates affecting growth stock valuations reducing positions and rising interest rates affecting growth stock valuations
""" """