From aff504c55c1ec49995450d98809872b713f7cc8f Mon Sep 17 00:00:00 2001 From: zoubobo Date: Thu, 23 Oct 2025 01:24:24 +0800 Subject: [PATCH] fix model call --- tradingagents/agents/utils/memory.py | 28 ++++++++++++++++++++++++---- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/tradingagents/agents/utils/memory.py b/tradingagents/agents/utils/memory.py index e1428204..7eced831 100644 --- a/tradingagents/agents/utils/memory.py +++ b/tradingagents/agents/utils/memory.py @@ -7,19 +7,39 @@ from sentence_transformers import SentenceTransformer class FinancialSituationMemory: def __init__(self, name, config): - if config["backend_url"] == "http://localhost:11434/v1": + self.config = config # Store config for backend routing + backend_url = self.config.get("backend_url", "") + + # Set embedding model based on backend + if "deepseek" in backend_url: + self.embedding = "deepseek-embed" + elif backend_url == "http://localhost:11434/v1": self.embedding = "nomic-embed-text" else: self.embedding = "text-embedding-3-small" + 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 OpenAI embedding for a text""" - if os.environ.get("OPENAI_API_KEY"): + """Get embedding for a text from the configured backend""" + backend_url = self.config.get("backend_url", "") + + if "deepseek" in backend_url: + client = OpenAI( + api_key=os.environ.get("DEEPSEEK_API_KEY"), + base_url=backend_url + ) + response = client.embeddings.create( + model=self.embedding, + input=text + ) + return response.data[0].embedding + elif os.environ.get("OPENAI_API_KEY") and "openai" in backend_url: client = OpenAI() response = client.embeddings.create( - model=self.embedding, input=text + model=self.embedding, + input=text ) return response.data[0].embedding else: