Jeremy Live
commited on
Commit
·
6380113
1
Parent(s):
332a246
upgrade config model gemini
Browse files
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...")
|