charieswei
/
TradingAgents
mirror of https://github.com/TauricResearch/TradingAgents.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

:[add] OpenAIEmbeddings from langchain_openai

This commit is contained in:
kimheesu 2025-06-12 15:06:35 +09:00
parent dd422893cf
commit 87713db002
43 changed files with 5230 additions and 5229 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
tradingagents/agents/utils/memory.py
View File

@ -2,20 +2,21 @@ import chromadb
from chromadb.config import Settings from chromadb.config import Settings
from openai import OpenAI from openai import OpenAI
import numpy as np import numpy as np
from langchain_openai import OpenAIEmbeddings
import os
class FinancialSituationMemory: class FinancialSituationMemory:
def __init__(self, name): def __init__(self, name):
self.client = OpenAI() # self.client = OpenAI()
self.embeddings = OpenAIEmbeddings(model="text-embedding-ada-002", api_key=os.getenv("OPENAI_API_KEY"))
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.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( embedding = self.embeddings.embed_query(text)
model="text-embedding-ada-002", input=text
) return embedding
return response.data[0].embedding
def add_situations(self, situations_and_advice): def add_situations(self, situations_and_advice):
"""Add financial situations and their corresponding advice. Parameter is a list of tuples (situation, rec)""" """Add financial situations and their corresponding advice. Parameter is a list of tuples (situation, rec)"""