Jeremy Live commited on
Commit
6380113
·
1 Parent(s): 332a246

upgrade config model gemini

Browse files
Files changed (1) hide show
  1. app.py +52 -2
app.py CHANGED
@@ -120,7 +120,8 @@ def initialize_llm():
120
  llm = ChatGoogleGenerativeAI(
121
  model="gemini-2.0-flash",
122
  temperature=0,
123
- google_api_key=google_api_key
 
124
  )
125
 
126
  # Test the model with a simple prompt
@@ -170,7 +171,6 @@ def create_agent():
170
  logger.error(error_msg)
171
  return None, error_msg
172
 
173
- # Step 3: Create SQL agent
174
  try:
175
  logger.info("Creating SQL agent...")
176
  agent = create_sql_agent(
@@ -196,6 +196,56 @@ def create_agent():
196
  logger.error(error_msg, exc_info=True) # Include full stack trace
197
  return None, error_msg
198
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
199
  # Inicializar el agente
200
  logger.info("="*50)
201
  logger.info("Starting application initialization...")
 
120
  llm = ChatGoogleGenerativeAI(
121
  model="gemini-2.0-flash",
122
  temperature=0,
123
+ google_api_key=google_api_key,
124
+ convert_system_message_to_human=True # Convert system messages to human messages
125
  )
126
 
127
  # Test the model with a simple prompt
 
171
  logger.error(error_msg)
172
  return None, error_msg
173
 
 
174
  try:
175
  logger.info("Creating SQL agent...")
176
  agent = create_sql_agent(
 
196
  logger.error(error_msg, exc_info=True) # Include full stack trace
197
  return None, error_msg
198
 
199
+ def create_agent(llm, db_connection):
200
+ """Create and return a SQL database agent."""
201
+ if not llm:
202
+ error_msg = "Cannot create agent: LLM is not available"
203
+ logger.error(error_msg)
204
+ return None, error_msg
205
+
206
+ if not db_connection:
207
+ error_msg = "Cannot create agent: Database connection is not available"
208
+ logger.error(error_msg)
209
+ return None, error_msg
210
+
211
+ try:
212
+ logger.info("Creating SQL agent...")
213
+
214
+ # Create the database toolkit with additional configuration
215
+ toolkit = SQLDatabaseToolkit(
216
+ db=db_connection,
217
+ llm=llm
218
+ )
219
+
220
+ # Create the agent with more detailed configuration
221
+ agent = create_sql_agent(
222
+ llm=llm,
223
+ toolkit=toolkit,
224
+ agent_type=AgentType.OPENAI_FUNCTIONS,
225
+ verbose=True,
226
+ handle_parsing_errors=True, # Better error handling for parsing
227
+ max_iterations=10, # Limit the number of iterations
228
+ early_stopping_method="generate" # Stop early if the agent is stuck
229
+ )
230
+
231
+ # Test the agent with a simple query
232
+ logger.info("Testing agent with a simple query...")
233
+ try:
234
+ test_query = "SELECT 1"
235
+ test_result = agent.run(test_query)
236
+ logger.info(f"Agent test query successful: {str(test_result)[:200]}...")
237
+ except Exception as e:
238
+ logger.warning(f"Agent test query failed (this might be expected): {str(e)}")
239
+ # Continue even if test fails, as it might be due to model limitations
240
+
241
+ logger.info("SQL agent created successfully")
242
+ return agent, ""
243
+
244
+ except Exception as e:
245
+ error_msg = f"Error creating SQL agent: {str(e)}"
246
+ logger.error(error_msg, exc_info=True)
247
+ return None, error_msg
248
+
249
  # Inicializar el agente
250
  logger.info("="*50)
251
  logger.info("Starting application initialization...")