Modify embedding model and requirements.txt

This commit is contained in:
Yuchen Zhang 2025-06-19 20:42:06 +08:00
parent 1d22ddae58
commit d7d7fc2a7d
4 changed files with 22 additions and 8 deletions

View File

@ -252,7 +252,8 @@ def select_llm_provider() -> tuple[str, str]:
("Anthropic", "https://api.anthropic.com/"),
("Google", "https://generativelanguage.googleapis.com/v1"),
("Openrouter", "https://openrouter.ai/api/v1"),
("Ollama", "http://localhost:11434/v1"),
("Ollama", "http://localhost:11434/v1"),
("DeepSeek", "https://api.deepseek.com/v1")
]
choice = questionary.select(

View File

@ -23,3 +23,7 @@ chainlit
rich
questionary
langchain_anthropic
langchain_deepseek
langchain_google_genai
openai
sentence-transformers

View File

@ -1,7 +1,9 @@
import os
import chromadb
from chromadb.config import Settings
from openai import OpenAI
from sentence_transformers import SentenceTransformer
class FinancialSituationMemory:
def __init__(self, name, config):
@ -9,17 +11,23 @@ class FinancialSituationMemory:
self.embedding = "nomic-embed-text"
else:
self.embedding = "text-embedding-3-small"
self.client = OpenAI()
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"""
response = self.client.embeddings.create(
model=self.embedding, input=text
)
return response.data[0].embedding
if os.environ.get("OPENAI_API_KEY"):
client = OpenAI()
response = client.embeddings.create(
model=self.embedding, input=text
)
return response.data[0].embedding
else:
model = SentenceTransformer("all-MiniLM-L6-v2")
embeddings = model.encode(text, convert_to_numpy=True, precision='int8')
# 单文本返回一维数组,多文本返回二维数组
return embeddings
def add_situations(self, situations_and_advice):
"""Add financial situations and their corresponding advice. Parameter is a list of tuples (situation, rec)"""

View File

@ -6,6 +6,7 @@ import json
from datetime import date
from typing import Dict, Any, Tuple, List, Optional
from langchain_deepseek import ChatDeepSeek
from langchain_openai import ChatOpenAI
from langchain_anthropic import ChatAnthropic
from langchain_google_genai import ChatGoogleGenerativeAI