pvanand commited on
Commit
4d0ec3e
1 Parent(s): 768bfff

Add full logging

Browse files
Files changed (1) hide show
  1. main.py +45 -1
main.py CHANGED
@@ -17,6 +17,18 @@ from prompts import CODING_ASSISTANT_PROMPT, NEWS_ASSISTANT_PROMPT, generate_new
17
  from fastapi_cache import FastAPICache
18
  from fastapi_cache.backends.inmemory import InMemoryBackend
19
  from fastapi_cache.decorator import cache
 
 
 
 
 
 
 
 
 
 
 
 
20
 
21
  app = FastAPI()
22
 
@@ -69,6 +81,7 @@ class NewsQueryModel(BaseModel):
69
 
70
  @lru_cache()
71
  def get_api_keys():
 
72
  return {
73
  "OPENROUTER_API_KEY": f"sk-or-v1-{os.environ['OPENROUTER_API_KEY']}",
74
  "BRAVE_API_KEY": os.environ['BRAVE_API_KEY']
@@ -91,6 +104,7 @@ def calculate_tokens(msgs):
91
  return sum(len(encoding.encode(str(m))) for m in msgs)
92
 
93
  def chat_with_llama_stream(messages, model="gpt-3.5-turbo", max_llm_history=4, max_output_tokens=2500):
 
94
  while calculate_tokens(messages) > (8000 - max_output_tokens):
95
  if len(messages) > max_llm_history:
96
  messages = [messages[0]] + messages[-max_llm_history:]
@@ -98,6 +112,7 @@ def chat_with_llama_stream(messages, model="gpt-3.5-turbo", max_llm_history=4, m
98
  max_llm_history -= 1
99
  if max_llm_history < 2:
100
  error_message = "Token limit exceeded. Please shorten your input or start a new conversation."
 
101
  raise HTTPException(status_code=400, detail=error_message)
102
 
103
  try:
@@ -117,11 +132,14 @@ def chat_with_llama_stream(messages, model="gpt-3.5-turbo", max_llm_history=4, m
117
 
118
  # After streaming, add the full response to the conversation history
119
  messages.append({"role": "assistant", "content": full_response})
 
120
  except Exception as e:
 
121
  raise HTTPException(status_code=500, detail=f"Error in model response: {str(e)}")
122
 
123
  async def verify_api_key(api_key: str = Security(api_key_header)):
124
  if api_key != API_KEY:
 
125
  raise HTTPException(status_code=403, detail="Could not validate credentials")
126
  return api_key
127
 
@@ -129,6 +147,7 @@ async def verify_api_key(api_key: str = Security(api_key_header)):
129
  DB_PATH = '/app/data/conversations.db'
130
 
131
  def init_db():
 
132
  os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
133
  conn = sqlite3.connect(DB_PATH)
134
  c = conn.cursor()
@@ -141,19 +160,23 @@ def init_db():
141
  timestamp DATETIME DEFAULT CURRENT_TIMESTAMP)''')
142
  conn.commit()
143
  conn.close()
 
144
 
145
  init_db()
146
 
147
  def update_db(user_id, conversation_id, message, response):
 
148
  conn = sqlite3.connect(DB_PATH)
149
  c = conn.cursor()
150
  c.execute('''INSERT INTO conversations (user_id, conversation_id, message, response)
151
  VALUES (?, ?, ?, ?)''', (user_id, conversation_id, message, response))
152
  conn.commit()
153
  conn.close()
 
154
 
155
  async def clear_inactive_conversations():
156
  while True:
 
157
  current_time = time.time()
158
  inactive_convos = [conv_id for conv_id, last_time in last_activity.items()
159
  if current_time - last_time > 1800] # 30 minutes
@@ -162,10 +185,12 @@ async def clear_inactive_conversations():
162
  del conversations[conv_id]
163
  if conv_id in last_activity:
164
  del last_activity[conv_id]
 
165
  await asyncio.sleep(60) # Check every minute
166
 
167
  @app.on_event("startup")
168
  async def startup_event():
 
169
  FastAPICache.init(InMemoryBackend(), prefix="fastapi-cache")
170
  asyncio.create_task(clear_inactive_conversations())
171
 
@@ -183,6 +208,7 @@ async def coding_assistant(query: QueryModel, background_tasks: BackgroundTasks,
183
  - google/gemma-2-27b-it
184
  Requires API Key authentication via X-API-Key header.
185
  """
 
186
  if query.conversation_id not in conversations:
187
  conversations[query.conversation_id] = [
188
  {"role": "system", "content": "You are a helpful assistant proficient in coding tasks. Help the user in understanding and writing code."}
@@ -200,12 +226,14 @@ async def coding_assistant(query: QueryModel, background_tasks: BackgroundTasks,
200
  full_response += content
201
  yield content
202
  background_tasks.add_task(update_db, query.user_id, query.conversation_id, query.user_query, full_response)
 
203
 
204
  return StreamingResponse(process_response(), media_type="text/event-stream")
205
 
206
  # New functions for news assistant
207
 
208
  def internet_search(query, type = "web", num_results=20):
 
209
  if type == "web":
210
  url = "https://api.search.brave.com/res/v1/web/search"
211
  else:
@@ -221,6 +249,7 @@ def internet_search(query, type = "web", num_results=20):
221
  response = requests.get(url, headers=headers, params=params)
222
 
223
  if response.status_code != 200:
 
224
  return []
225
 
226
  if type == "web":
@@ -240,17 +269,21 @@ def internet_search(query, type = "web", num_results=20):
240
  }
241
  processed_results.append(result)
242
 
 
243
  return processed_results[:num_results]
244
 
245
  @lru_cache(maxsize=100)
246
  def cached_internet_search(query: str):
 
247
  return internet_search(query, type = "news")
248
 
249
 
250
  def analyze_news(query):
 
251
  news_data = cached_internet_search(query)
252
 
253
  if not news_data:
 
254
  return "Failed to fetch news data.", []
255
 
256
  # Prepare the prompt for the AI
@@ -262,6 +295,7 @@ def analyze_news(query):
262
  {"role": "user", "content": prompt}
263
  ]
264
 
 
265
  return messages
266
 
267
  @app.post("/news-assistant")
@@ -270,15 +304,18 @@ async def news_assistant(query: NewsQueryModel, api_key: str = Depends(verify_ap
270
  News assistant endpoint that provides summaries and analysis of recent news based on user queries.
271
  Requires API Key authentication via X-API-Key header.
272
  """
 
273
  messages = analyze_news(query.query)
274
 
275
  if not messages:
 
276
  raise HTTPException(status_code=500, detail="Failed to fetch news data")
277
 
278
  def process_response():
279
  for content in chat_with_llama_stream(messages, model=query.model_id):
280
  yield content
281
- #meta-llama/llama-3-70b-instruct google/gemini-pro-1.5
 
282
  return StreamingResponse(process_response(), media_type="text/event-stream")
283
 
284
  class SearchQueryModel(BaseModel):
@@ -296,9 +333,11 @@ class SearchQueryModel(BaseModel):
296
  }
297
 
298
  def analyze_search_results(query):
 
299
  search_data = internet_search(query, type="web")
300
 
301
  if not search_data:
 
302
  return "Failed to fetch search data.", []
303
 
304
  # Prepare the prompt for the AI
@@ -309,6 +348,7 @@ def analyze_search_results(query):
309
  {"role": "user", "content": prompt}
310
  ]
311
 
 
312
  return messages
313
 
314
  @app.post("/search-assistant")
@@ -317,17 +357,21 @@ async def search_assistant(query: SearchQueryModel, api_key: str = Depends(verif
317
  Search assistant endpoint that provides summaries and analysis of web search results based on user queries.
318
  Requires API Key authentication via X-API-Key header.
319
  """
 
320
  messages = analyze_search_results(query.query)
321
 
322
  if not messages:
 
323
  raise HTTPException(status_code=500, detail="Failed to fetch search data")
324
 
325
  def process_response():
326
  for content in chat_with_llama_stream(messages, model=query.model_id):
327
  yield content
 
328
 
329
  return StreamingResponse(process_response(), media_type="text/event-stream")
330
 
331
  if __name__ == "__main__":
332
  import uvicorn
 
333
  uvicorn.run(app, host="0.0.0.0", port=7860)
 
17
  from fastapi_cache import FastAPICache
18
  from fastapi_cache.backends.inmemory import InMemoryBackend
19
  from fastapi_cache.decorator import cache
20
+ import logging
21
+
22
+ # Configure logging
23
+ logging.basicConfig(
24
+ level=logging.INFO,
25
+ format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
26
+ handlers=[
27
+ logging.FileHandler("app.log"),
28
+ logging.StreamHandler()
29
+ ]
30
+ )
31
+ logger = logging.getLogger(__name__)
32
 
33
  app = FastAPI()
34
 
 
81
 
82
  @lru_cache()
83
  def get_api_keys():
84
+ logger.info("Loading API keys")
85
  return {
86
  "OPENROUTER_API_KEY": f"sk-or-v1-{os.environ['OPENROUTER_API_KEY']}",
87
  "BRAVE_API_KEY": os.environ['BRAVE_API_KEY']
 
104
  return sum(len(encoding.encode(str(m))) for m in msgs)
105
 
106
  def chat_with_llama_stream(messages, model="gpt-3.5-turbo", max_llm_history=4, max_output_tokens=2500):
107
+ logger.info(f"Starting chat with model: {model}")
108
  while calculate_tokens(messages) > (8000 - max_output_tokens):
109
  if len(messages) > max_llm_history:
110
  messages = [messages[0]] + messages[-max_llm_history:]
 
112
  max_llm_history -= 1
113
  if max_llm_history < 2:
114
  error_message = "Token limit exceeded. Please shorten your input or start a new conversation."
115
+ logger.error(error_message)
116
  raise HTTPException(status_code=400, detail=error_message)
117
 
118
  try:
 
132
 
133
  # After streaming, add the full response to the conversation history
134
  messages.append({"role": "assistant", "content": full_response})
135
+ logger.info("Chat completed successfully")
136
  except Exception as e:
137
+ logger.error(f"Error in model response: {str(e)}")
138
  raise HTTPException(status_code=500, detail=f"Error in model response: {str(e)}")
139
 
140
  async def verify_api_key(api_key: str = Security(api_key_header)):
141
  if api_key != API_KEY:
142
+ logger.warning("Invalid API key used")
143
  raise HTTPException(status_code=403, detail="Could not validate credentials")
144
  return api_key
145
 
 
147
  DB_PATH = '/app/data/conversations.db'
148
 
149
  def init_db():
150
+ logger.info("Initializing database")
151
  os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
152
  conn = sqlite3.connect(DB_PATH)
153
  c = conn.cursor()
 
160
  timestamp DATETIME DEFAULT CURRENT_TIMESTAMP)''')
161
  conn.commit()
162
  conn.close()
163
+ logger.info("Database initialized successfully")
164
 
165
  init_db()
166
 
167
  def update_db(user_id, conversation_id, message, response):
168
+ logger.info(f"Updating database for conversation: {conversation_id}")
169
  conn = sqlite3.connect(DB_PATH)
170
  c = conn.cursor()
171
  c.execute('''INSERT INTO conversations (user_id, conversation_id, message, response)
172
  VALUES (?, ?, ?, ?)''', (user_id, conversation_id, message, response))
173
  conn.commit()
174
  conn.close()
175
+ logger.info("Database updated successfully")
176
 
177
  async def clear_inactive_conversations():
178
  while True:
179
+ logger.info("Clearing inactive conversations")
180
  current_time = time.time()
181
  inactive_convos = [conv_id for conv_id, last_time in last_activity.items()
182
  if current_time - last_time > 1800] # 30 minutes
 
185
  del conversations[conv_id]
186
  if conv_id in last_activity:
187
  del last_activity[conv_id]
188
+ logger.info(f"Cleared {len(inactive_convos)} inactive conversations")
189
  await asyncio.sleep(60) # Check every minute
190
 
191
  @app.on_event("startup")
192
  async def startup_event():
193
+ logger.info("Starting up the application")
194
  FastAPICache.init(InMemoryBackend(), prefix="fastapi-cache")
195
  asyncio.create_task(clear_inactive_conversations())
196
 
 
208
  - google/gemma-2-27b-it
209
  Requires API Key authentication via X-API-Key header.
210
  """
211
+ logger.info(f"Received coding assistant query: {query.user_query}")
212
  if query.conversation_id not in conversations:
213
  conversations[query.conversation_id] = [
214
  {"role": "system", "content": "You are a helpful assistant proficient in coding tasks. Help the user in understanding and writing code."}
 
226
  full_response += content
227
  yield content
228
  background_tasks.add_task(update_db, query.user_id, query.conversation_id, query.user_query, full_response)
229
+ logger.info(f"Completed coding assistant response for query: {query.user_query}")
230
 
231
  return StreamingResponse(process_response(), media_type="text/event-stream")
232
 
233
  # New functions for news assistant
234
 
235
  def internet_search(query, type = "web", num_results=20):
236
+ logger.info(f"Performing internet search for query: {query}, type: {type}")
237
  if type == "web":
238
  url = "https://api.search.brave.com/res/v1/web/search"
239
  else:
 
249
  response = requests.get(url, headers=headers, params=params)
250
 
251
  if response.status_code != 200:
252
+ logger.error(f"Failed to fetch search results. Status code: {response.status_code}")
253
  return []
254
 
255
  if type == "web":
 
269
  }
270
  processed_results.append(result)
271
 
272
+ logger.info(f"Retrieved {len(processed_results)} search results")
273
  return processed_results[:num_results]
274
 
275
  @lru_cache(maxsize=100)
276
  def cached_internet_search(query: str):
277
+ logger.info(f"Performing cached internet search for query: {query}")
278
  return internet_search(query, type = "news")
279
 
280
 
281
  def analyze_news(query):
282
+ logger.info(f"Analyzing news for query: {query}")
283
  news_data = cached_internet_search(query)
284
 
285
  if not news_data:
286
+ logger.error("Failed to fetch news data")
287
  return "Failed to fetch news data.", []
288
 
289
  # Prepare the prompt for the AI
 
295
  {"role": "user", "content": prompt}
296
  ]
297
 
298
+ logger.info("News analysis completed")
299
  return messages
300
 
301
  @app.post("/news-assistant")
 
304
  News assistant endpoint that provides summaries and analysis of recent news based on user queries.
305
  Requires API Key authentication via X-API-Key header.
306
  """
307
+ logger.info(f"Received news assistant query: {query.query}")
308
  messages = analyze_news(query.query)
309
 
310
  if not messages:
311
+ logger.error("Failed to fetch news data")
312
  raise HTTPException(status_code=500, detail="Failed to fetch news data")
313
 
314
  def process_response():
315
  for content in chat_with_llama_stream(messages, model=query.model_id):
316
  yield content
317
+ logger.info(f"Completed news assistant response for query: {query.query}")
318
+
319
  return StreamingResponse(process_response(), media_type="text/event-stream")
320
 
321
  class SearchQueryModel(BaseModel):
 
333
  }
334
 
335
  def analyze_search_results(query):
336
+ logger.info(f"Analyzing search results for query: {query}")
337
  search_data = internet_search(query, type="web")
338
 
339
  if not search_data:
340
+ logger.error("Failed to fetch search data")
341
  return "Failed to fetch search data.", []
342
 
343
  # Prepare the prompt for the AI
 
348
  {"role": "user", "content": prompt}
349
  ]
350
 
351
+ logger.info("Search results analysis completed")
352
  return messages
353
 
354
  @app.post("/search-assistant")
 
357
  Search assistant endpoint that provides summaries and analysis of web search results based on user queries.
358
  Requires API Key authentication via X-API-Key header.
359
  """
360
+ logger.info(f"Received search assistant query: {query.query}")
361
  messages = analyze_search_results(query.query)
362
 
363
  if not messages:
364
+ logger.error("Failed to fetch search data")
365
  raise HTTPException(status_code=500, detail="Failed to fetch search data")
366
 
367
  def process_response():
368
  for content in chat_with_llama_stream(messages, model=query.model_id):
369
  yield content
370
+ logger.info(f"Completed search assistant response for query: {query.query}")
371
 
372
  return StreamingResponse(process_response(), media_type="text/event-stream")
373
 
374
  if __name__ == "__main__":
375
  import uvicorn
376
+ logger.info("Starting the application")
377
  uvicorn.run(app, host="0.0.0.0", port=7860)