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/"), ("Anthropic", "https://api.anthropic.com/"),
("Google", "https://generativelanguage.googleapis.com/v1"), ("Google", "https://generativelanguage.googleapis.com/v1"),
("Openrouter", "https://openrouter.ai/api/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( choice = questionary.select(

View File

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

View File

@ -1,7 +1,9 @@
import os
import chromadb import chromadb
from chromadb.config import Settings from chromadb.config import Settings
from openai import OpenAI from openai import OpenAI
from sentence_transformers import SentenceTransformer
class FinancialSituationMemory: class FinancialSituationMemory:
def __init__(self, name, config): def __init__(self, name, config):
@ -9,17 +11,23 @@ class FinancialSituationMemory:
self.embedding = "nomic-embed-text" self.embedding = "nomic-embed-text"
else: else:
self.embedding = "text-embedding-3-small" self.embedding = "text-embedding-3-small"
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.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"""
if os.environ.get("OPENAI_API_KEY"):
response = self.client.embeddings.create( client = OpenAI()
model=self.embedding, input=text response = client.embeddings.create(
) model=self.embedding, input=text
return response.data[0].embedding )
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): 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)"""

View File

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