pvanand commited on
Commit
2f5381c
·
verified ·
1 Parent(s): 50f8031

Delete document_generator_v2.py

Browse files
Files changed (1) hide show
  1. document_generator_v2.py +0 -725
document_generator_v2.py DELETED
@@ -1,725 +0,0 @@
1
- # File: prompts.py
2
-
3
- DOCUMENT_OUTLINE_PROMPT_SYSTEM = """You are a document generator. Provide the outline of the document requested in <prompt></prompt> in JSON format.
4
- Include sections and subsections if required. Use the "Content" field to provide a specific prompt or instruction for generating content for that particular section or subsection.
5
- make sure the Sections follow a logical flow and each prompt's content does not overlap with other sections.
6
- OUTPUT IN FOLLOWING JSON FORMAT enclosed in <output> tags
7
- <output>
8
- {
9
- "Document": {
10
- "Title": "Document Title",
11
- "Author": "Author Name",
12
- "Date": "YYYY-MM-DD",
13
- "Version": "1.0",
14
-
15
- "Sections": [
16
- {
17
- "SectionNumber": "1",
18
- "Title": "Section Title",
19
- "Content": "Specific prompt or instruction for generating content for this section",
20
- "Subsections": [
21
- {
22
- "SectionNumber": "1.1",
23
- "Title": "Subsection Title",
24
- "Content": "Specific prompt or instruction for generating content for this subsection"
25
- }
26
- ]
27
- }
28
- ]
29
- }
30
- }
31
- </output>"""
32
-
33
- DOCUMENT_OUTLINE_PROMPT_USER = """<prompt>{query}</prompt>"""
34
-
35
- DOCUMENT_SECTION_PROMPT_SYSTEM = """You are a document generator, You need to output only the content requested in the section in the prompt.
36
- FORMAT YOUR OUTPUT AS MARKDOWN ENCLOSED IN <response></response> tags
37
- <overall_objective>{overall_objective}</overall_objective>
38
- <document_layout>{document_layout}</document_layout>"""
39
-
40
- DOCUMENT_SECTION_PROMPT_USER = """<prompt>Output the content for the section "{section_or_subsection_title}" formatted as markdown. Follow this instruction: {content_instruction}</prompt>"""
41
-
42
- ##########################################
43
-
44
- DOCUMENT_TEMPLATE_OUTLINE_PROMPT_SYSTEM = """You are a document template generator. Provide the outline of the document requested in <prompt></prompt> in JSON format.
45
- Include sections and subsections if required. Use the "Content" field to provide a specific prompt or instruction for generating template with placeholder text /example content for that particular section or subsection. Specify in each prompt to output as a template and use placeholder text/ tables as necessory.
46
- make sure the Sections follow a logical flow and each prompt's content does not overlap with other sections.
47
- OUTPUT IN FOLLOWING JSON FORMAT enclosed in <output> tags
48
- <output>
49
- {
50
- "Document": {
51
- "Title": "Document Title",
52
- "Author": "Author Name",
53
- "Date": "YYYY-MM-DD",
54
- "Version": "1.0",
55
-
56
- "Sections": [
57
- {
58
- "SectionNumber": "1",
59
- "Title": "Section Title",
60
- "Content": "Specific prompt or instruction for generating template for this section",
61
- "Subsections": [
62
- {
63
- "SectionNumber": "1.1",
64
- "Title": "Subsection Title",
65
- "Content": "Specific prompt or instruction for generating template for this subsection"
66
- }
67
- ]
68
- }
69
- ]
70
- }
71
- }
72
- </output>"""
73
-
74
- DOCUMENT_TEMPLATE_PROMPT_USER = """<prompt>{query}</prompt>"""
75
-
76
- DOCUMENT_TEMPLATE_SECTION_PROMPT_SYSTEM = """You are a document template generator,You need to output only the content requested in the section in the prompt, Use placeholder text/examples/tables wherever required.
77
- FORMAT YOUR OUTPUT AS A TEMPLATE ENCLOSED IN <response></response> tags
78
- <overall_objective>{overall_objective}</overall_objective>
79
- <document_layout>{document_layout}</document_layout>"""
80
-
81
- DOCUMENT_TEMPLATE_SECTION_PROMPT_USER = """<prompt>Output the content for the section "{section_or_subsection_title}" formatted as markdown. Follow this instruction: {content_instruction}</prompt>"""
82
-
83
-
84
- # File: llm_observability.py
85
-
86
- import sqlite3
87
- import json
88
- from datetime import datetime
89
- from typing import Dict, Any, List, Optional
90
-
91
- class LLMObservabilityManager:
92
- def __init__(self, db_path: str = "llm_observability_v2.db"):
93
- self.db_path = db_path
94
- self.create_table()
95
-
96
- def create_table(self):
97
- with sqlite3.connect(self.db_path) as conn:
98
- cursor = conn.cursor()
99
- cursor.execute('''
100
- CREATE TABLE IF NOT EXISTS llm_observations (
101
- id TEXT PRIMARY KEY,
102
- conversation_id TEXT,
103
- created_at DATETIME,
104
- status TEXT,
105
- request TEXT,
106
- response TEXT,
107
- model TEXT,
108
- total_tokens INTEGER,
109
- prompt_tokens INTEGER,
110
- completion_tokens INTEGER,
111
- latency FLOAT,
112
- user TEXT
113
- )
114
- ''')
115
-
116
- def insert_observation(self, response: Dict[str, Any], conversation_id: str, status: str, request: str, latency: float, user: str):
117
- created_at = datetime.fromtimestamp(response['created'])
118
-
119
- with sqlite3.connect(self.db_path) as conn:
120
- cursor = conn.cursor()
121
- cursor.execute('''
122
- INSERT INTO llm_observations
123
- (id, conversation_id, created_at, status, request, response, model, total_tokens, prompt_tokens, completion_tokens, latency, user)
124
- VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
125
- ''', (
126
- response['id'],
127
- conversation_id,
128
- created_at,
129
- status,
130
- request,
131
- json.dumps(response['choices'][0]['message']),
132
- response['model'],
133
- response['usage']['total_tokens'],
134
- response['usage']['prompt_tokens'],
135
- response['usage']['completion_tokens'],
136
- latency,
137
- user
138
- ))
139
-
140
- def get_observations(self, conversation_id: Optional[str] = None) -> List[Dict[str, Any]]:
141
- with sqlite3.connect(self.db_path) as conn:
142
- cursor = conn.cursor()
143
- if conversation_id:
144
- cursor.execute('SELECT * FROM llm_observations WHERE conversation_id = ? ORDER BY created_at', (conversation_id,))
145
- else:
146
- cursor.execute('SELECT * FROM llm_observations ORDER BY created_at')
147
- rows = cursor.fetchall()
148
-
149
- column_names = [description[0] for description in cursor.description]
150
- return [dict(zip(column_names, row)) for row in rows]
151
-
152
- def get_all_observations(self) -> List[Dict[str, Any]]:
153
- return self.get_observations()
154
-
155
-
156
- # File: app.py
157
- import os
158
- import json
159
- import re
160
- import asyncio
161
- import time
162
- from typing import List, Dict, Optional, Any, Callable, Union
163
- from openai import OpenAI
164
- import logging
165
- import functools
166
- from fastapi import APIRouter, HTTPException, Request, UploadFile, File, Depends
167
- from fastapi.responses import StreamingResponse
168
- from pydantic import BaseModel
169
- from fastapi_cache import FastAPICache
170
- from fastapi_cache.decorator import cache
171
- import psycopg2
172
- from datetime import datetime
173
- import base64
174
- from fastapi import Form
175
- from llama_parse import LlamaParse
176
-
177
- logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
178
- logger = logging.getLogger(__name__)
179
-
180
- def log_execution(func: Callable) -> Callable:
181
- @functools.wraps(func)
182
- def wrapper(*args: Any, **kwargs: Any) -> Any:
183
- logger.info(f"Executing {func.__name__}")
184
- try:
185
- result = func(*args, **kwargs)
186
- logger.info(f"{func.__name__} completed successfully")
187
- return result
188
- except Exception as e:
189
- logger.error(f"Error in {func.__name__}: {e}")
190
- raise
191
- return wrapper
192
-
193
- # aiclient.py
194
-
195
- class AIClient:
196
- def __init__(self):
197
- self.client = OpenAI(
198
- base_url="https://openrouter.ai/api/v1",
199
- api_key="sk-or-v1-" + os.environ['OPENROUTER_API_KEY']
200
- )
201
- self.observability_manager = LLMObservabilityManager()
202
-
203
- @log_execution
204
- def generate_response(
205
- self,
206
- messages: List[Dict[str, str]],
207
- model: str = "openai/gpt-4o-mini",
208
- max_tokens: int = 32000,
209
- conversation_id: str = None,
210
- user: str = "anonymous"
211
- ) -> Optional[str]:
212
- if not messages:
213
- return None
214
-
215
- start_time = time.time()
216
- response = self.client.chat.completions.create(
217
- model=model,
218
- messages=messages,
219
- max_tokens=max_tokens,
220
- stream=False
221
- )
222
- end_time = time.time()
223
- latency = end_time - start_time
224
-
225
- # Log the observation
226
- self.observability_manager.insert_observation(
227
- response=response.dict(),
228
- conversation_id=conversation_id or "default",
229
- status="success",
230
- request=json.dumps(messages),
231
- latency=latency,
232
- user=user
233
- )
234
-
235
- return response.choices[0].message.content
236
-
237
- @log_execution
238
- def generate_vision_response(
239
- self,
240
- messages: List[Dict[str, Union[str, List[Dict[str, Union[str, Dict[str, str]]]]]]],
241
- model: str = "google/gemini-flash-1.5-8b",
242
- max_tokens: int = 32000,
243
- conversation_id: str = None,
244
- user: str = "anonymous"
245
- ) -> Optional[str]:
246
- if not messages:
247
- return None
248
-
249
- start_time = time.time()
250
- response = self.client.chat.completions.create(
251
- model=model,
252
- messages=messages,
253
- max_tokens=max_tokens,
254
- stream=False
255
- )
256
- end_time = time.time()
257
- latency = end_time - start_time
258
-
259
- # Log the observation
260
- self.observability_manager.insert_observation(
261
- response=response.dict(),
262
- conversation_id=conversation_id or "default",
263
- status="success",
264
- request=json.dumps(messages),
265
- latency=latency,
266
- user=user
267
- )
268
-
269
- return response.choices[0].message.content
270
-
271
-
272
- class VisionTools:
273
- def __init__(self, ai_client):
274
- self.ai_client = ai_client
275
-
276
- async def extract_images_info(self, images: List[UploadFile]) -> str:
277
- try:
278
- image_contents = []
279
- for image in images:
280
- image_content = await image.read()
281
- base64_image = base64.b64encode(image_content).decode('utf-8')
282
- image_contents.append({
283
- "type": "image_url",
284
- "image_url": {
285
- "url": f"data:image/jpeg;base64,{base64_image}"
286
- }
287
- })
288
-
289
- messages = [
290
- {
291
- "role": "user",
292
- "content": [
293
- {
294
- "type": "text",
295
- "text": "Extract the contents of these images in detail in a structured format, focusing on any text, tables, diagrams, or visual elements that might be relevant for document generation."
296
- },
297
- *image_contents
298
- ]
299
- }
300
- ]
301
-
302
- image_context = self.ai_client.generate_vision_response(messages)
303
- return image_context
304
- except Exception as e:
305
- print(f"Error processing images: {str(e)}")
306
- return ""
307
-
308
-
309
- class DatabaseManager:
310
- """Manages database operations."""
311
-
312
- def __init__(self):
313
- self.db_params = {
314
- "dbname": "postgres",
315
- "user": os.environ['SUPABASE_USER'],
316
- "password": os.environ['SUPABASE_PASSWORD'],
317
- "host": "aws-0-us-west-1.pooler.supabase.com",
318
- "port": "5432"
319
- }
320
-
321
- @log_execution
322
- def update_database(self, user_id: str, user_query: str, response: str) -> None:
323
- with psycopg2.connect(**self.db_params) as conn:
324
- with conn.cursor() as cur:
325
- insert_query = """
326
- INSERT INTO ai_document_generator (user_id, user_query, response)
327
- VALUES (%s, %s, %s);
328
- """
329
- cur.execute(insert_query, (user_id, user_query, response))
330
-
331
- class DocumentGenerator:
332
- def __init__(self, ai_client: AIClient):
333
- self.ai_client = ai_client
334
- self.document_outline = None
335
- self.content_messages = []
336
-
337
- @staticmethod
338
- def extract_between_tags(text: str, tag: str) -> str:
339
- pattern = f"<{tag}>(.*?)</{tag}>"
340
- match = re.search(pattern, text, re.DOTALL)
341
- return match.group(1).strip() if match else ""
342
-
343
- @staticmethod
344
- def remove_duplicate_title(content: str, title: str, section_number: str) -> str:
345
- patterns = [
346
- rf"^#+\s*{re.escape(section_number)}(?:\s+|\s*:\s*|\.\s*){re.escape(title)}",
347
- rf"^#+\s*{re.escape(title)}",
348
- rf"^{re.escape(section_number)}(?:\s+|\s*:\s*|\.\s*){re.escape(title)}",
349
- rf"^{re.escape(title)}",
350
- ]
351
-
352
- for pattern in patterns:
353
- content = re.sub(pattern, "", content, flags=re.MULTILINE | re.IGNORECASE)
354
-
355
- return content.lstrip()
356
-
357
- @log_execution
358
- def generate_document_outline(self, query: str, template: bool = False, image_context: str = "", max_retries: int = 3) -> Optional[Dict]:
359
- messages = [
360
- {"role": "system", "content": DOCUMENT_OUTLINE_PROMPT_SYSTEM if not template else DOCUMENT_TEMPLATE_OUTLINE_PROMPT_SYSTEM},
361
- {"role": "user", "content": DOCUMENT_OUTLINE_PROMPT_USER.format(query=query) if not template else DOCUMENT_TEMPLATE_PROMPT_USER.format(query=query, image_context=image_context)}
362
- ]
363
- # Update user content to include image context if provided
364
- if image_context:
365
- messages[1]["content"] += f"<attached_images>\n\n{image_context}\n\n</attached_images>"
366
-
367
- for attempt in range(max_retries):
368
- outline_response = self.ai_client.generate_response(messages, model="openai/gpt-4o")
369
- outline_json_text = self.extract_between_tags(outline_response, "output")
370
-
371
- try:
372
- self.document_outline = json.loads(outline_json_text)
373
- return self.document_outline
374
- except json.JSONDecodeError as e:
375
- if attempt < max_retries - 1:
376
- logger.warning(f"Failed to parse JSON (attempt {attempt + 1}): {e}")
377
- logger.info("Retrying...")
378
- else:
379
- logger.error(f"Failed to parse JSON after {max_retries} attempts: {e}")
380
- return None
381
-
382
- @log_execution
383
- def generate_content(self, title: str, content_instruction: str, section_number: str, template: bool = False) -> str:
384
- SECTION_PROMPT_USER = DOCUMENT_SECTION_PROMPT_USER if not template else DOCUMENT_TEMPLATE_SECTION_PROMPT_USER
385
- self.content_messages.append({
386
- "role": "user",
387
- "content": SECTION_PROMPT_USER.format(
388
- section_or_subsection_title=title,
389
- content_instruction=content_instruction
390
- )
391
- })
392
- section_response = self.ai_client.generate_response(self.content_messages)
393
- content = self.extract_between_tags(section_response, "response")
394
- content = self.remove_duplicate_title(content, title, section_number)
395
- self.content_messages.append({
396
- "role": "assistant",
397
- "content": section_response
398
- })
399
- return content
400
-
401
- class MarkdownConverter:
402
- @staticmethod
403
- def slugify(text: str) -> str:
404
- return re.sub(r'\W+', '-', text.lower())
405
-
406
- @classmethod
407
- def generate_toc(cls, sections: List[Dict]) -> str:
408
- toc = "<div style='page-break-before: always;'></div>\n\n"
409
- toc += "<h2 style='color: #2c3e50; text-align: center;'>Table of Contents</h2>\n\n"
410
- toc += "<nav style='background-color: #f8f9fa; padding: 20px; border-radius: 5px; line-height: 1.6;'>\n\n"
411
- for section in sections:
412
- section_number = section['SectionNumber']
413
- section_title = section['Title']
414
- toc += f"<p><a href='#{cls.slugify(section_title)}' style='color: #3498db; text-decoration: none;'>{section_number}. {section_title}</a></p>\n\n"
415
-
416
- for subsection in section.get('Subsections', []):
417
- subsection_number = subsection['SectionNumber']
418
- subsection_title = subsection['Title']
419
- toc += f"<p style='margin-left: 20px;'><a href='#{cls.slugify(subsection_title)}' style='color: #2980b9; text-decoration: none;'>{subsection_number} {subsection_title}</a></p>\n\n"
420
-
421
- toc += "</nav>\n\n"
422
- return toc
423
-
424
- @classmethod
425
- def convert_to_markdown(cls, document: Dict) -> str:
426
- markdown = "<div style='text-align: center; padding-top: 33vh;'>\n\n"
427
- markdown += f"<h1 style='color: #2c3e50; border-bottom: 2px solid #3498db; padding-bottom: 10px; display: inline-block;'>{document['Title']}</h1>\n\n"
428
- markdown += f"<p style='color: #7f8c8d;'><em>By {document['Author']}</em></p>\n\n"
429
- markdown += f"<p style='color: #95a5a6;'>Version {document['Version']} | {document['Date']}</p>\n\n"
430
- markdown += "</div>\n\n"
431
-
432
- markdown += cls.generate_toc(document['Sections'])
433
-
434
- markdown += "<div style='max-width: 800px; margin: 0 auto; font-family: \"Segoe UI\", Arial, sans-serif; line-height: 1.6;'>\n\n"
435
-
436
- for section in document['Sections']:
437
- markdown += "<div style='page-break-before: always;'></div>\n\n"
438
- section_number = section['SectionNumber']
439
- section_title = section['Title']
440
- markdown += f"<h2 id='{cls.slugify(section_title)}' style='color: #2c3e50; border-bottom: 1px solid #bdc3c7; padding-bottom: 5px;'>{section_number}. {section_title}</h2>\n\n"
441
- markdown += f"<div style='color: #34495e; margin-bottom: 20px;'>\n\n{section['Content']}\n\n</div>\n\n"
442
-
443
- for subsection in section.get('Subsections', []):
444
- subsection_number = subsection['SectionNumber']
445
- subsection_title = subsection['Title']
446
- markdown += f"<h3 id='{cls.slugify(subsection_title)}' style='color: #34495e;'>{subsection_number} {subsection_title}</h3>\n\n"
447
- markdown += f"<div style='color: #34495e; margin-bottom: 20px;'>\n\n{subsection['Content']}\n\n</div>\n\n"
448
-
449
- markdown += "</div>"
450
- return markdown
451
-
452
- async def load_documents(documents: List[UploadFile]) -> List[str]:
453
- """
454
- Load and parse documents using LlamaParse.
455
-
456
- Args:
457
- documents (List[UploadFile]): List of uploaded document files.
458
-
459
- Returns:
460
- List[str]: List of parsed document contents.
461
- """
462
- parser = LlamaParse(
463
- api_key=os.getenv("LLAMA_PARSE_API_KEY"),
464
- result_type="markdown",
465
- num_workers=4,
466
- verbose=True,
467
- language="en",
468
- )
469
-
470
- # Save uploaded files temporarily
471
- temp_files = []
472
- for doc in documents:
473
- temp_file_path = f"/tmp/{doc.filename}"
474
- with open(temp_file_path, "wb") as buffer:
475
- content = await doc.read()
476
- buffer.write(content)
477
- temp_files.append(temp_file_path)
478
-
479
- try:
480
- # Use LlamaParse to extract content
481
- print(f"processing files {str(temp_files)}")
482
- parsed_documents = await parser.aload_data(temp_files)
483
- documents_list = [doc.text for doc in parsed_documents]
484
- return documents_list
485
- finally:
486
- # Clean up temporary files
487
- for temp_file in temp_files:
488
- os.remove(temp_file)
489
-
490
-
491
-
492
- router = APIRouter()
493
-
494
- class JsonDocumentResponse(BaseModel):
495
- json_document: Dict
496
-
497
- # class JsonDocumentRequest(BaseModel):
498
- # query: str
499
- # template: bool = False
500
- # images: Optional[List[UploadFile]] = File(None)
501
- # documents: Optional[List[UploadFile]] = File(None)
502
- # conversation_id: str = ""
503
-
504
- class MarkdownDocumentRequest(BaseModel):
505
- json_document: Dict
506
- query: str
507
- template: bool = False
508
- conversation_id: str = ""
509
-
510
- MESSAGE_DELIMITER = b"\n---DELIMITER---\n"
511
-
512
- def yield_message(message):
513
- message_json = json.dumps(message, ensure_ascii=False).encode('utf-8')
514
- return message_json + MESSAGE_DELIMITER
515
-
516
- async def generate_document_stream(document_generator: DocumentGenerator, document_outline: Dict, query: str, template: bool = False, conversation_id: str = ""):
517
- document_generator.document_outline = document_outline
518
- db_manager = DatabaseManager()
519
- overall_objective = query
520
- document_layout = json.dumps(document_generator.document_outline, indent=2)
521
- cache_key = f"image_context_{conversation_id}"
522
- image_context = await FastAPICache.get_backend().get(cache_key)
523
-
524
- SECTION_PROMPT_SYSTEM = DOCUMENT_SECTION_PROMPT_SYSTEM if not template else DOCUMENT_TEMPLATE_SECTION_PROMPT_SYSTEM
525
- document_generator.content_messages = [
526
- {
527
- "role": "system",
528
- "content": SECTION_PROMPT_SYSTEM.format(
529
- overall_objective=overall_objective,
530
- document_layout=document_layout
531
- )
532
- }
533
- ]
534
- if image_context:
535
- document_generator.content_messages[0]["content"] += f"<attached_images>\n\n{image_context}\n\n</attached_images>"
536
-
537
- for section in document_generator.document_outline["Document"].get("Sections", []):
538
- section_title = section.get("Title", "")
539
- section_number = section.get("SectionNumber", "")
540
- content_instruction = section.get("Content", "")
541
- logging.info(f"Generating content for section: {section_title}")
542
- content = document_generator.generate_content(section_title, content_instruction, section_number, template)
543
- section["Content"] = content
544
- yield yield_message({
545
- "type": "document_section",
546
- "content": {
547
- "section_number": section_number,
548
- "section_title": section_title,
549
- "content": content
550
- }
551
- })
552
-
553
- for subsection in section.get("Subsections", []):
554
- subsection_title = subsection.get("Title", "")
555
- subsection_number = subsection.get("SectionNumber", "")
556
- subsection_content_instruction = subsection.get("Content", "")
557
- logging.info(f"Generating content for subsection: {subsection_title}")
558
- content = document_generator.generate_content(subsection_title, subsection_content_instruction, subsection_number, template)
559
- subsection["Content"] = content
560
- yield yield_message({
561
- "type": "document_section",
562
- "content": {
563
- "section_number": subsection_number,
564
- "section_title": subsection_title,
565
- "content": content
566
- }
567
- })
568
-
569
- markdown_document = MarkdownConverter.convert_to_markdown(document_generator.document_outline["Document"])
570
-
571
- yield yield_message({
572
- "type": "complete_document",
573
- "content": {
574
- "markdown": markdown_document,
575
- "json": document_generator.document_outline
576
- },
577
- });
578
-
579
- db_manager.update_database("elevatics", query, markdown_document)
580
-
581
- @router.post("/generate-document/markdown-stream")
582
- async def generate_markdown_document_stream_endpoint(request: MarkdownDocumentRequest):
583
- ai_client = AIClient()
584
- document_generator = DocumentGenerator(ai_client)
585
-
586
- async def stream_generator():
587
- try:
588
- async for chunk in generate_document_stream(document_generator, request.json_document, request.query, request.template, request.conversation_id):
589
- yield chunk
590
- except Exception as e:
591
- yield yield_message({
592
- "type": "error",
593
- "content": str(e)
594
- })
595
-
596
- return StreamingResponse(stream_generator(), media_type="application/octet-stream")
597
-
598
-
599
- @cache(expire=600*24*7)
600
- @router.post("/generate-document/json", response_model=JsonDocumentResponse)
601
- async def generate_document_outline_endpoint(
602
- query: str = Form(...),
603
- template: bool = Form(False),
604
- conversation_id: str = Form(...),
605
- images: Optional[List[UploadFile]] = File(None),
606
- documents: Optional[List[UploadFile]] = File(None)
607
- ):
608
- ai_client = AIClient()
609
- document_generator = DocumentGenerator(ai_client)
610
- vision_tools = VisionTools(ai_client)
611
-
612
- try:
613
- # Handle image processing
614
- image_context = ""
615
- if images:
616
- image_context = await vision_tools.extract_images_info(images)
617
- # Store the image_context in the cache
618
- image_cache_key = f"image_context_{conversation_id}"
619
- await FastAPICache.get_backend().set(image_cache_key, image_context, expire=3600) # Cache for 1 hour
620
-
621
- # Handle document processing using the new load_documents function
622
- documents_list = []
623
- if documents:
624
- documents_list = await load_documents(documents)
625
- # Store the documents_list in the cache
626
- docs_cache_key = f"documents_list_{conversation_id}"
627
- print("saving document as cache key:",docs_cache_key)
628
- await FastAPICache.get_backend().set(docs_cache_key, documents_list, expire=3600) # Cache for 1 hour
629
-
630
- # Generate document outline
631
- json_document = document_generator.generate_document_outline(
632
- query,
633
- template,
634
- image_context=image_context,
635
- #documents_context=documents_list
636
- )
637
-
638
- if json_document is None:
639
- raise HTTPException(status_code=500, detail="Failed to generate a valid document outline")
640
-
641
- return JsonDocumentResponse(json_document=json_document)
642
-
643
- except Exception as e:
644
- raise HTTPException(status_code=500, detail=str(e))
645
-
646
-
647
- ## OBSERVABILITY
648
- from uuid import uuid4
649
- import csv
650
- from io import StringIO
651
-
652
- class ObservationResponse(BaseModel):
653
- observations: List[Dict]
654
-
655
- def create_csv_response(observations: List[Dict]) -> StreamingResponse:
656
- def iter_csv(data):
657
- output = StringIO()
658
- writer = csv.DictWriter(output, fieldnames=data[0].keys() if data else [])
659
- writer.writeheader()
660
- for row in data:
661
- writer.writerow(row)
662
- output.seek(0)
663
- yield output.read()
664
-
665
- headers = {
666
- 'Content-Disposition': 'attachment; filename="observations.csv"'
667
- }
668
- return StreamingResponse(iter_csv(observations), media_type="text/csv", headers=headers)
669
-
670
-
671
- @router.get("/last-observations/{limit}")
672
- async def get_last_observations(limit: int = 10, format: str = "json"):
673
- observability_manager = LLMObservabilityManager()
674
-
675
- try:
676
- # Get all observations, sorted by created_at in descending order
677
- all_observations = observability_manager.get_observations()
678
- all_observations.sort(key=lambda x: x['created_at'], reverse=True)
679
-
680
- # Get the last conversation_id
681
- if all_observations:
682
- last_conversation_id = all_observations[0]['conversation_id']
683
-
684
- # Filter observations for the last conversation
685
- last_conversation_observations = [
686
- obs for obs in all_observations
687
- if obs['conversation_id'] == last_conversation_id
688
- ][:limit]
689
-
690
- if format.lower() == "csv":
691
- return create_csv_response(last_conversation_observations)
692
- else:
693
- return ObservationResponse(observations=last_conversation_observations)
694
- else:
695
- if format.lower() == "csv":
696
- return create_csv_response([])
697
- else:
698
- return ObservationResponse(observations=[])
699
- except Exception as e:
700
- raise HTTPException(status_code=500, detail=f"Failed to retrieve observations: {str(e)}")
701
-
702
- ## TEST CACHE
703
-
704
- class CacheItem(BaseModel):
705
- key: str
706
- value: str
707
-
708
- @router.post("/set-cache")
709
- async def set_cache(item: CacheItem):
710
- try:
711
- # Set the cache with a default expiration of 1 hour (3600 seconds)
712
- await FastAPICache.get_backend().set(item.key, item.value, expire=3600)
713
- return {"message": f"Cache set for key: {item.key}"}
714
- except Exception as e:
715
- raise HTTPException(status_code=500, detail=f"Failed to set cache: {str(e)}")
716
-
717
- @router.get("/get-cache/{key}")
718
- async def get_cache(key: str):
719
- try:
720
- value = await FastAPICache.get_backend().get(key)
721
- if value is None:
722
- raise HTTPException(status_code=404, detail=f"No cache found for key: {key}")
723
- return {"key": key, "value": value}
724
- except Exception as e:
725
- raise HTTPException(status_code=500, detail=f"Failed to get cache: {str(e)}")