dygoo commited on
Commit
4d6ad08
·
verified ·
1 Parent(s): c1ec90a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +41 -5
app.py CHANGED
@@ -5,12 +5,17 @@ import inspect
5
  import pandas as pd
6
  import smolagents
7
  from smolagents import DuckDuckGoSearchTool, VisitWebpageTool
8
- import time # Added: For sleep functionality
9
- from functools import lru_cache # Added: For caching search results
 
 
 
10
 
11
  from youtube_transcript_api import YouTubeTranscriptApi
12
  import re
13
 
 
 
14
 
15
  class YouTubeVideoTool:
16
  def __init__(self):
@@ -155,10 +160,41 @@ class BasicAgent:
155
 
156
 
157
 
158
- def _formulate_direct_answer(self, relevant_info, question):
159
- if self.model:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
160
  return f"Based on the search: {relevant_info}"
161
- return relevant_info
 
162
 
163
 
164
 
 
5
  import pandas as pd
6
  import smolagents
7
  from smolagents import DuckDuckGoSearchTool, VisitWebpageTool
8
+ import time
9
+ from functools import lru_cache
10
+
11
+ import google.generativeai as genai
12
+ from google.generativeai.types import HarmCategory, HarmBlockThreshold
13
 
14
  from youtube_transcript_api import YouTubeTranscriptApi
15
  import re
16
 
17
+ GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
18
+ genai.configure(api_key=GOOGLE_API_KEY)
19
 
20
  class YouTubeVideoTool:
21
  def __init__(self):
 
160
 
161
 
162
 
163
+ def _formulate_direct_answer(self, relevant_info, question):
164
+ if self.model and self.model.startswith('gemini'):
165
+ try:
166
+ # Configure the model
167
+ generation_config = {
168
+ "temperature": 0.7,
169
+ "top_p": 0.95,
170
+ "top_k": 40,
171
+ "max_output_tokens": 1024,
172
+ }
173
+
174
+ safety_settings = {
175
+ HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
176
+ HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
177
+ HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
178
+ HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
179
+ }
180
+
181
+ # Initialize the model
182
+ model = genai.GenerativeModel(
183
+ model_name="gemini-pro", # Adjust as needed based on your model string
184
+ generation_config=generation_config,
185
+ safety_settings=safety_settings
186
+ )
187
+
188
+ # Prepare prompt and generate response
189
+ prompt = f"Question: {question}\n\nRelevant information: {relevant_info}\n\nProvide a concise answer based only on the given information."
190
+ response = model.generate_content(prompt)
191
+ return response.text
192
+
193
+ except Exception as e:
194
+ print(f"Error using Gemini model: {e}")
195
  return f"Based on the search: {relevant_info}"
196
+
197
+ return relevant_info
198
 
199
 
200