deepdive-IR / agent /adapters.py
Ritabanm's picture
Upload 25 files
94b06be verified
raw
history blame contribute delete
692 Bytes
from typing import List, Dict
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from settings import OPENAI_API_KEY
class LLMAdapter:
def __init__(self, model: str = "gpt-4o-mini", temperature: float = 0.2):
self.llm = ChatOpenAI(model=model, temperature=temperature, api_key=OPENAI_API_KEY)
def chat(self, messages: List[Dict[str, str]]) -> str:
return self.llm.invoke(messages).content
class EmbeddingsAdapter:
def __init__(self, model: str = "text-embedding-3-small"):
self.emb = OpenAIEmbeddings(model=model, api_key=OPENAI_API_KEY)
def embed(self, texts: List[str]) -> List[List[float]]:
return self.emb.embed_documents(texts)