Added support to use openrouter models
Gaurav Dhiman committed
Aug 3, 2024 at 23:30 UTC
cb60dd4fbd3c25ea3c1aedcd91259884a77cf076
3 files changed
+10
example.env
+1
@@ -3,6 +3,7 @@ API_KEY_ANTHROPIC=
3
API_KEY_GROQ=
4
API_KEY_PERPLEXITY=
5
API_KEY_GOOGLE=
6
+API_KEY_OPENROUTER=
7
8
TOKENIZERS_PARALLELISM=true
9
PYDEVD_DISABLE_FILE_VALIDATION=1
\ No newline at end of file
main.py
+1
@@ -31,6 +31,7 @@ def initialize():
31
# chat_llm = models.get_ollama(model_name="llama3:8b-text-fp16")
32
# chat_llm = models.get_ollama(model_name="gemma2:latest")
33
# chat_llm = models.get_ollama(model_name="qwen:14b")
34
+ # chat_llm = models.get_openrouter(model_name="meta-llama/llama-3-8b-instruct:free")
35
# chat_llm = models.get_google_chat()
36
37
models.py
+8
@@ -6,6 +6,7 @@ from langchain_anthropic import ChatAnthropic
6
from langchain_groq import ChatGroq
7
from langchain_huggingface import HuggingFaceEmbeddings
8
from langchain_google_genai import ChatGoogleGenerativeAI, HarmBlockThreshold, HarmCategory
9
+from pydantic.v1.types import SecretStr
10
11
12
# Load environment variables
@@ -92,6 +93,13 @@ def get_google_chat(model_name="gemini-1.5-flash-latest", api_key=None, temperat
93
return ChatGoogleGenerativeAI(model=model_name, temperature=temperature, google_api_key=api_key,
94
safety_settings={HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE }) # type: ignore
95
96
+def get_openrouter(model_name: str="meta-llama/llama-3.1-8b-instruct:free"):
97
+ open_router_api_key = os.getenv('API_KEY_OPENROUTER')
98
+ open_router_api_key = SecretStr(open_router_api_key) if open_router_api_key else None
99
+ return ChatOpenAI(api_key=open_router_api_key,
100
+ base_url="https://openrouter.ai/api/v1",
101
+ model=model_name)
102
+
103
def get_embedding_hf(model_name="sentence-transformers/all-MiniLM-L6-v2"):
104
return HuggingFaceEmbeddings(model_name=model_name)
105