llamasrock commited on
Commit
fc680c0
·
verified ·
1 Parent(s): c6fcb72

Update app.py

Browse files

Changed the basic agent to Gemini client

Files changed (1) hide show
  1. app.py +24 -8
app.py CHANGED
@@ -18,20 +18,36 @@ GEMINI_API_KEY = os.getenv("Gemini_API_key")
18
 
19
  class BasicAgent:
20
  def __init__(self):
21
- self.llm_model = LiteLLMModel(
22
- model_id="gemini/gemini-2.0-flash", # you can see other model names here: https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models. It is important to prefix the name with "gemini/"
23
- api_key=GEMINI_API_KEY,
24
- max_tokens=8192
25
- )
26
- self.google_search_tool = Tool(google_search = GoogleSearch())
27
- self.agent = CodeAgent(model = self.llm_model, tools = [self.google_search_tool])
 
 
 
 
 
28
  print("BasicAgent initialized.")
29
  def __call__(self, question: str) -> str:
30
  print(f"Agent received question (first 50 chars): {question[:50]}...")
31
  fixed_answer = "This is a default answer."
32
  # print(f"Agent returning fixed answer: {fixed_answer}")
33
  # return fixed_answer
34
- return self.agent.run(question)
 
 
 
 
 
 
 
 
 
 
 
35
 
36
  def run_and_submit_all( profile: gr.OAuthProfile | None):
37
  """
 
18
 
19
  class BasicAgent:
20
  def __init__(self):
21
+ # self.llm_model = LiteLLMModel(
22
+ # model_id="gemini/gemini-2.0-flash", # you can see other model names here: https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models. It is important to prefix the name with "gemini/"
23
+ # api_key=GEMINI_API_KEY,
24
+ # max_tokens=8192
25
+ # )
26
+ # # self.google_search_tool = Tool(google_search = GoogleSearch())
27
+ # self.google_search_tool = GoogleSearchTool()
28
+ # self.agent = CodeAgent(model = self.llm_model, tools = [self.google_search_tool])
29
+ # Define Google API client with GoogleSearch tool
30
+ self.client = genai.Client(api_key="GEMINI_API_KEY")
31
+
32
+
33
  print("BasicAgent initialized.")
34
  def __call__(self, question: str) -> str:
35
  print(f"Agent received question (first 50 chars): {question[:50]}...")
36
  fixed_answer = "This is a default answer."
37
  # print(f"Agent returning fixed answer: {fixed_answer}")
38
  # return fixed_answer
39
+ # return self.agent.run(question)
40
+ # Use the Google GenAI client to run the question
41
+ response = self.client.models.generate_content(
42
+ model='gemini-1.5-flash',
43
+ contents=question,
44
+ config=types.GenerateContentConfig(
45
+ tools=[types.Tool(
46
+ google_search_retrieval=types.GoogleSearchRetrieval()
47
+ )]
48
+ )
49
+ )
50
+ return response
51
 
52
  def run_and_submit_all( profile: gr.OAuthProfile | None):
53
  """