diff --git a/.env.example b/.env.example index 1e257c3c..68ad0d51 100644 --- a/.env.example +++ b/.env.example @@ -1,2 +1,7 @@ ALPHA_VANTAGE_API_KEY=alpha_vantage_api_key_placeholder -OPENAI_API_KEY=openai_api_key_placeholder \ No newline at end of file +# only select 1 to your .env +OPENAI_API_KEY=openai_api_key_placeholder +OPENROUTER_API_KEY=openrouter_api_key_placeholder +ANTHROPIC_API_KEY=anthropic_api_key_placeholder +GOOGLE_API_KEY=google_api_key_placeholder +XAI_API_KEY=xai_api_key_placeholder \ No newline at end of file diff --git a/tradingagents/agents/utils/memory.py b/tradingagents/agents/utils/memory.py index ce9fdf11..4ace3290 100644 --- a/tradingagents/agents/utils/memory.py +++ b/tradingagents/agents/utils/memory.py @@ -6,8 +6,6 @@ from sentence_transformers import SentenceTransformer class FinancialSituationMemory: def __init__(self, name, config): - self.client = OpenAI(base_url=config["backend_url"], api_key=config.get("api_key")) - # Based on llm_provider to select embeddings if config["llm_provider"] == "ollama": self.embedding = "nomic-embed-text" @@ -18,14 +16,14 @@ class FinancialSituationMemory: self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2') self.use_local_embedding = True else: + self.client = OpenAI(base_url=config["backend_url"], api_key=config.get("api_key")) self.embedding = "text-embedding-3-small" self.embedding_client = self.client self.use_local_embedding = False - + self.chroma_client = chromadb.Client(Settings(allow_reset=True)) self.situation_collection = self.chroma_client.create_collection(name=name) - def get_embedding(self, text): """Get embedding for a text""" if self.use_local_embedding: @@ -34,12 +32,11 @@ class FinancialSituationMemory: else: # use API response = self.embedding_client.embeddings.create( - model=self.embedding, + model=self.embedding, input=text ) return response.data[0].embedding - def add_situations(self, situations_and_advice): """Add financial situations and their corresponding advice. Parameter is a list of tuples (situation, rec)"""