33 lines
1.1 KiB
Python
33 lines
1.1 KiB
Python
"""
|
|
LLM integration module for LangChain with support for OpenAI and Ollama.
|
|
Provides a reusable LLM client based on environment configuration.
|
|
"""
|
|
|
|
import os
|
|
from typing import Union
|
|
|
|
from langchain.llms import OpenAI, Ollama
|
|
from langchain.chat_models import ChatOpenAI, ChatOllama
|
|
|
|
# Environment variable to select provider: "openai" or "ollama"
|
|
LLM_PROVIDER = os.getenv("LLM_PROVIDER", "openai").lower()
|
|
|
|
|
|
def get_llm() -> Union[OpenAI, Ollama, ChatOpenAI, ChatOllama]:
|
|
"""
|
|
Returns an LLM instance based on the configured provider.
|
|
|
|
For OpenAI, uses the default OpenAI LLM (text-davinci-003 or gpt-3.5-turbo).
|
|
For Ollama, uses the default Ollama LLM (e.g., llama2).
|
|
|
|
Raises:
|
|
ValueError: If an unsupported provider is specified.
|
|
"""
|
|
if LLM_PROVIDER == "openai":
|
|
# Use ChatOpenAI for GPT-3.5-turbo by default
|
|
return ChatOpenAI(temperature=0.7)
|
|
elif LLM_PROVIDER == "ollama":
|
|
# Use ChatOllama for local models
|
|
return ChatOllama(model="llama2", temperature=0.7)
|
|
else:
|
|
raise ValueError(f"Unsupported LLM provider: {LLM_PROVIDER}") |