iQuentin commited on
Commit
0d66232
·
verified ·
1 Parent(s): b5f47fa

Add OpenAIServerModel

Browse files
Files changed (1) hide show
  1. agent.py +53 -17
agent.py CHANGED
@@ -5,7 +5,7 @@ from smolagents import (
5
  VisitWebpageTool,
6
  PythonInterpreterTool,
7
  InferenceClientModel,
8
- HfApiModel, # import bug from smolagents after adding duckduckgo-search in requirements
9
  tool
10
  )
11
  from typing import List, Dict, Any, Optional
@@ -31,22 +31,22 @@ class QAgent:
31
  self.system_prompt = system_prompt
32
 
33
 
34
- if model_type == "HfApiModel":
35
- if api_key is None:
36
- api_key = os.getenv("SP_HF_TOK")
37
- if not api_key:
38
- raise ValueError("No API Key found for HuggingFace. Please set SP_HF_TOK or pass api_key.")
39
-
40
- if self.verbose:
41
- print(f"Using Hugging Face token: {api_key[:5]}... (HfApiModel mode)")
42
-
43
- self.model = HfApiModel(
44
- model_id=model_id or "Qwen/Qwen2.5-Coder-32B-Instruct", # précédemment : or "meta-llama/Llama-3-70B-Instruct",
45
- token=api_key
46
- # temperature=temperature
47
- )
48
-
49
- elif model_type == "InferenceClientModel":
50
  if api_key is None:
51
  api_key = os.getenv("Q_NEB_TOK")
52
  if not api_key:
@@ -62,6 +62,42 @@ class QAgent:
62
  timeout=timeout or 120
63
  # temperature=temperature
64
  )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
65
  else:
66
  raise ValueError(f"Unknown model type: {model_type}")
67
 
 
5
  VisitWebpageTool,
6
  PythonInterpreterTool,
7
  InferenceClientModel,
8
+ # HfApiModel, # import bug from smolagents after adding duckduckgo-search in requirements
9
  tool
10
  )
11
  from typing import List, Dict, Any, Optional
 
31
  self.system_prompt = system_prompt
32
 
33
 
34
+ # if model_type == "HfApiModel":
35
+ # if api_key is None:
36
+ # api_key = os.getenv("Q_NEB_TOK")
37
+ # if not api_key:
38
+ # raise ValueError("No API Key found for HuggingFace. Please set Q_NEB_TOK or pass api_key.")
39
+ #
40
+ # if self.verbose:
41
+ # print(f"Using Hugging Face token: {api_key[:5]}... (HfApiModel mode)")
42
+ #
43
+ # self.model = HfApiModel(
44
+ # model_id=model_id or "Qwen/Qwen2.5-Coder-32B-Instruct", # précédemment : or "meta-llama/Llama-3-70B-Instruct",
45
+ # token=api_key
46
+ # # temperature=temperature
47
+ # )
48
+ # el
49
+ if model_type == "InferenceClientModel":
50
  if api_key is None:
51
  api_key = os.getenv("Q_NEB_TOK")
52
  if not api_key:
 
62
  timeout=timeout or 120
63
  # temperature=temperature
64
  )
65
+ elif model_type == "OpenAIServerModel":
66
+ print(f"Trying to configure OpenAIServerModel.")
67
+ # Check for xAI API key and base URL first
68
+ xai_api_key = os.getenv("XAI_API_KEY")
69
+ # xai_api_base = os.getenv("XAI_API_BASE") # Ne sais pas à quoi ça sert..
70
+
71
+ # If xAI credentials are available, use them
72
+ if xai_api_key and api_key is None:
73
+ api_key = xai_api_key
74
+ if self.verbose:
75
+ print(f"Using xAI API key: {api_key[:5]}...")
76
+
77
+ # If no API key specified, fall back to OPENAI_API_KEY
78
+ if api_key is None:
79
+ api_key = os.getenv("Q_OAI_TOK")
80
+ if not api_key:
81
+ raise ValueError("No OpenAI API key provided. Please set Q_OAI_TOK or XAI_API_KEY environment variable or pass api_key parameter.")
82
+
83
+ # If xAI API base is available and no api_base is provided, use it
84
+ if xai_api_base and api_base is None:
85
+ api_base = xai_api_base
86
+ if self.verbose:
87
+ print(f"Using xAI API base URL: {api_base}")
88
+
89
+ # If no API base specified but environment variable available, use it
90
+ if api_base is None:
91
+ api_base = os.getenv("AGENT_API_BASE")
92
+ if api_base and self.verbose:
93
+ print(f"Using API base from AGENT_API_BASE: {api_base}")
94
+
95
+ self.model = OpenAIServerModel(
96
+ model_id=model_id or "gpt-4o",
97
+ api_key=api_key,
98
+ # api_base=api_base,
99
+ temperature=temperature
100
+ )
101
  else:
102
  raise ValueError(f"Unknown model type: {model_type}")
103