sivan22 commited on
Commit
50cb379
·
verified ·
1 Parent(s): 0a496a7

Update llm_providers.py

Browse files
Files changed (1) hide show
  1. llm_providers.py +76 -77
llm_providers.py CHANGED
@@ -1,77 +1,76 @@
1
- from langchain_anthropic import ChatAnthropic
2
- from langchain_openai import ChatOpenAI
3
- from langchain_ollama import ChatOllama
4
- from langchain_core.language_models.base import BaseLanguageModel
5
- from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
6
- from typing import Optional, Dict, List, Any, Sequence
7
- from langchain.tools import BaseTool
8
- import os
9
- import requests
10
- import json
11
- from dotenv import load_dotenv
12
- from dataclasses import dataclass
13
- import ollama
14
- import copy
15
- from chat_gemini import ChatGemini
16
-
17
-
18
- load_dotenv()
19
-
20
-
21
-
22
- class LLMProvider:
23
-
24
- def __init__(self, api_keys: Optional[Dict[str, str]] = None):
25
- if api_keys:
26
- self.api_keys = api_keys
27
- self.providers: Dict[str, Any] = {}
28
- self._setup_providers()
29
-
30
- def _get_ollama_models(self) -> List[str]:
31
- """Get list of available Ollama models using the ollama package"""
32
- try:
33
- models = ollama.list()
34
- return [model.model for model in models['models']]
35
- except Exception:
36
- return []
37
-
38
- def _setup_providers(self):
39
- os.environ['REQUESTS_CA_BUNDLE'] = 'C:\\ProgramData\\NetFree\\CA\\netfree-ca-bundle-curl.crt'
40
-
41
- # Google Gemini
42
- if google_key := os.getenv('GOOGLE_API_KEY') or self.api_keys.get('google'):
43
- self.providers['Gemini'] = ChatGemini(api_key=google_key)
44
-
45
- # Anthropicsel
46
- if anthropic_key := os.getenv('ANTHROPIC_API_KEY') or self.api_keys.get('anthropic'):
47
- self.providers['Claude'] = ChatAnthropic(
48
- api_key=anthropic_key,
49
- model_name="claude-3-5-sonnet-20241022",
50
- )
51
-
52
- # OpenAI
53
- if openai_key := os.getenv('OPENAI_API_KEY') or self.api_keys.get('openai'):
54
- self.providers['ChatGPT'] = ChatOpenAI(
55
- api_key=openai_key,
56
- model_name="gpt-4o-2024-11-20",
57
- max_completion_tokens=4096,
58
-
59
- )
60
-
61
- # Ollama (local)
62
- try:
63
- # Get available Ollama models using the ollama package
64
- ollama_models = self._get_ollama_models()
65
- for model in ollama_models:
66
- self.providers[f'Ollama-{model}'] = ChatOllama(model=model)
67
- except Exception:
68
- pass # Ollama not available
69
-
70
- def get_available_providers(self) -> list[str]:
71
- """Return list of available provider names"""
72
-
73
- return list(self.providers.keys())
74
-
75
- def get_provider(self, name: str) -> Optional[Any]:
76
- """Get LLM provider by name"""
77
- return self.providers.get(name)
 
1
+ from langchain_anthropic import ChatAnthropic
2
+ from langchain_openai import ChatOpenAI
3
+ from langchain_ollama import ChatOllama
4
+ from langchain_core.language_models.base import BaseLanguageModel
5
+ from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
6
+ from typing import Optional, Dict, List, Any, Sequence
7
+ from langchain.tools import BaseTool
8
+ import os
9
+ import requests
10
+ import json
11
+ from dotenv import load_dotenv
12
+ from dataclasses import dataclass
13
+ import ollama
14
+ import copy
15
+ from chat_gemini import ChatGemini
16
+
17
+
18
+ load_dotenv()
19
+
20
+
21
+
22
+ class LLMProvider:
23
+
24
+ def __init__(self, api_keys: Optional[Dict[str, str]] = None):
25
+ if api_keys:
26
+ self.api_keys = api_keys
27
+ self.providers: Dict[str, Any] = {}
28
+ self._setup_providers()
29
+
30
+ def _get_ollama_models(self) -> List[str]:
31
+ """Get list of available Ollama models using the ollama package"""
32
+ try:
33
+ models = ollama.list()
34
+ return [model.model for model in models['models']]
35
+ except Exception:
36
+ return []
37
+
38
+ def _setup_providers(self):
39
+
40
+ # Google Gemini
41
+ if google_key := os.getenv('GOOGLE_API_KEY') or self.api_keys.get('google'):
42
+ self.providers['Gemini'] = ChatGemini(api_key=google_key)
43
+
44
+ # Anthropicsel
45
+ if anthropic_key := os.getenv('ANTHROPIC_API_KEY') or self.api_keys.get('anthropic'):
46
+ self.providers['Claude'] = ChatAnthropic(
47
+ api_key=anthropic_key,
48
+ model_name="claude-3-5-sonnet-20241022",
49
+ )
50
+
51
+ # OpenAI
52
+ if openai_key := os.getenv('OPENAI_API_KEY') or self.api_keys.get('openai'):
53
+ self.providers['ChatGPT'] = ChatOpenAI(
54
+ api_key=openai_key,
55
+ model_name="gpt-4o-2024-11-20",
56
+ max_completion_tokens=4096,
57
+
58
+ )
59
+
60
+ # Ollama (local)
61
+ try:
62
+ # Get available Ollama models using the ollama package
63
+ ollama_models = self._get_ollama_models()
64
+ for model in ollama_models:
65
+ self.providers[f'Ollama-{model}'] = ChatOllama(model=model)
66
+ except Exception:
67
+ pass # Ollama not available
68
+
69
+ def get_available_providers(self) -> list[str]:
70
+ """Return list of available provider names"""
71
+
72
+ return list(self.providers.keys())
73
+
74
+ def get_provider(self, name: str) -> Optional[Any]:
75
+ """Get LLM provider by name"""
76
+ return self.providers.get(name)