File size: 12,209 Bytes
0289fc8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ada64d9
457ce08
0289fc8
 
 
 
 
 
e81c7e9
0289fc8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ff71d34
0289fc8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
443d394
0289fc8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e3cfbf6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
# File: prompts.py

DOCUMENT_OUTLINE_PROMPT_SYSTEM = """You are a document generator. Provide the outline of the document requested in <prompt></prompt> in JSON format.
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.

OUTPUT IN FOLLOWING JSON FORMAT enclosed in <output> tags
<output>
{
    "Document": {
      "Title": "Document Title",
      "Author": "Author Name",
      "Date": "YYYY-MM-DD",
      "Version": "1.0",

      "Sections": [
        {
          "SectionNumber": "1",
          "Title": "Section Title",
          "Content": "Specific prompt or instruction for generating content for this section",
          "Subsections": [
            {
              "SectionNumber": "1.1",
              "Title": "Subsection Title",
              "Content": "Specific prompt or instruction for generating content for this subsection"
            }
          ]
        }
      ]
    }
  }
</output>"""

DOCUMENT_OUTLINE_PROMPT_USER = """<prompt>{query}</prompt>"""

DOCUMENT_SECTION_PROMPT_SYSTEM = """You are a document generator, You need to output only the content requested in the section in the prompt.
FORMAT YOUR OUTPUT AS MARKDOWN ENCLOSED IN <response></response> tags
<overall_objective>{overall_objective}</overall_objective>
<document_layout>{document_layout}</document_layout>"""

DOCUMENT_SECTION_PROMPT_USER = """<prompt>Output the content for the section "{section_or_subsection_title}" formatted as markdown. Follow this instruction: {content_instruction}</prompt>"""

# File: app.py

import os
import json
import re
import time
import asyncio
from typing import List, Dict, Optional, Any, Callable
from openai import OpenAI
import logging
import functools
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from fastapi_cache.decorator import cache
#from prompts import *

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def log_execution(func: Callable) -> Callable:
    @functools.wraps(func)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        logger.info(f"Executing {func.__name__}")
        try:
            result = func(*args, **kwargs)
            logger.info(f"{func.__name__} completed successfully")
            return result
        except Exception as e:
            logger.error(f"Error in {func.__name__}: {e}")
            raise
    return wrapper

class AIClient:
    def __init__(self):
        self.client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key="sk-or-v1-"+os.environ['OPENROUTER_API_KEY']
        )
    @log_execution
    def generate_response(
        self,
        messages: List[Dict[str, str]],
        model: str = "openai/gpt-4o-mini",
        max_tokens: int = 32000
    ) -> Optional[str]:
        if not messages:
            return None
        response = self.client.chat.completions.create(
            model=model,
            messages=messages,
            max_tokens=max_tokens,
            stream=False
        )
        return response.choices[0].message.content

class DocumentGenerator:
    def __init__(self, ai_client: AIClient):
        self.ai_client = ai_client
        self.document_outline = None
        self.content_messages = []

    @staticmethod
    def extract_between_tags(text: str, tag: str) -> str:
        pattern = f"<{tag}>(.*?)</{tag}>"
        match = re.search(pattern, text, re.DOTALL)
        return match.group(1).strip() if match else ""

    @staticmethod
    def remove_duplicate_title(content: str, title: str, section_number: str) -> str:
        patterns = [
            rf"^#+\s*{re.escape(section_number)}(?:\s+|\s*:\s*|\.\s*){re.escape(title)}",
            rf"^#+\s*{re.escape(title)}",
            rf"^{re.escape(section_number)}(?:\s+|\s*:\s*|\.\s*){re.escape(title)}",
            rf"^{re.escape(title)}",
        ]
        
        for pattern in patterns:
            content = re.sub(pattern, "", content, flags=re.MULTILINE | re.IGNORECASE)
        
        return content.lstrip()

    @log_execution
    def generate_document_outline(self, query: str, max_retries: int = 3) -> Optional[Dict]:
        messages = [
            {"role": "system", "content": DOCUMENT_OUTLINE_PROMPT_SYSTEM},
            {"role": "user", "content": DOCUMENT_OUTLINE_PROMPT_USER.format(query=query)}
        ]
        
        for attempt in range(max_retries):
            outline_response = self.ai_client.generate_response(messages, model="openai/gpt-4o")
            outline_json_text = self.extract_between_tags(outline_response, "output")
            
            try:
                self.document_outline = json.loads(outline_json_text)
                return self.document_outline
            except json.JSONDecodeError as e:
                if attempt < max_retries - 1:
                    logger.warning(f"Failed to parse JSON (attempt {attempt + 1}): {e}")
                    logger.info("Retrying...")
                else:
                    logger.error(f"Failed to parse JSON after {max_retries} attempts: {e}")
                    return None

    @log_execution
    def generate_content(self, title: str, content_instruction: str, section_number: str) -> str:
        self.content_messages.append({
            "role": "user",
            "content": DOCUMENT_SECTION_PROMPT_USER.format(
                section_or_subsection_title=title,
                content_instruction=content_instruction
            )
        })
        section_response = self.ai_client.generate_response(self.content_messages)
        content = self.extract_between_tags(section_response, "response")
        content = self.remove_duplicate_title(content, title, section_number)
        self.content_messages.append({
            "role": "assistant",
            "content": section_response
        })
        return content

    @log_execution
    def generate_document(self, query: str) -> Dict:
        self.generate_document_outline(query)
        
        if self.document_outline is None:
            raise ValueError("Failed to generate a valid document outline")

        overall_objective = query
        document_layout = json.dumps(self.document_outline, indent=2)

        self.content_messages = [
            {
                "role": "system",
                "content": DOCUMENT_SECTION_PROMPT_SYSTEM.format(
                    overall_objective=overall_objective,
                    document_layout=document_layout
                )
            }
        ]

        for section in self.document_outline["Document"].get("Sections", []):
            section_title = section.get("Title", "")
            section_number = section.get("SectionNumber", "")
            content_instruction = section.get("Content", "")
            logger.info(f"Generating content for section: {section_title}")
            section["Content"] = self.generate_content(section_title, content_instruction, section_number)

            for subsection in section.get("Subsections", []):
                subsection_title = subsection.get("Title", "")
                subsection_number = subsection.get("SectionNumber", "")
                subsection_content_instruction = subsection.get("Content", "")
                logger.info(f"Generating content for subsection: {subsection_title}")
                subsection["Content"] = self.generate_content(subsection_title, subsection_content_instruction, subsection_number)

        return self.document_outline


class MarkdownConverter:
    @staticmethod
    def slugify(text: str) -> str:
        return re.sub(r'\W+', '-', text.lower())

    @classmethod
    def generate_toc(cls, sections: List[Dict]) -> str:
        toc = "<div style='page-break-before: always;'></div>\n\n"
        toc += "<h2 style='color: #2c3e50; text-align: center;'>Table of Contents</h2>\n\n"
        toc += "<nav style='background-color: #f8f9fa; padding: 20px; border-radius: 5px; line-height: 1.6;'>\n\n"
        for section in sections:
            section_number = section['SectionNumber']
            section_title = section['Title']
            toc += f"<p><a href='#{cls.slugify(section_title)}' style='color: #3498db; text-decoration: none;'>{section_number}. {section_title}</a></p>\n\n"
            
            for subsection in section.get('Subsections', []):
                subsection_number = subsection['SectionNumber']
                subsection_title = subsection['Title']
                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"
        
        toc += "</nav>\n\n"
        return toc

    @classmethod
    def convert_to_markdown(cls, document: Dict) -> str:
        # First page with centered content
        markdown = "<div style='text-align: center; padding-top: 33vh;'>\n\n"
        markdown += f"<h1 style='color: #2c3e50; border-bottom: 2px solid #3498db; padding-bottom: 10px; display: inline-block;'>{document['Title']}</h1>\n\n"
        markdown += f"<p style='color: #7f8c8d;'><em>By {document['Author']}</em></p>\n\n"
        markdown += f"<p style='color: #95a5a6;'>Version {document['Version']} | {document['Date']}</p>\n\n"
        markdown += "</div>\n\n"
        
        # Table of Contents on the second page
        markdown += cls.generate_toc(document['Sections'])
        
        # Main content
        markdown += "<div style='max-width: 800px; margin: 0 auto; font-family: \"Segoe UI\", Arial, sans-serif; line-height: 1.6;'>\n\n"
        
        for section in document['Sections']:
            markdown += "<div style='page-break-before: always;'></div>\n\n"
            section_number = section['SectionNumber']
            section_title = section['Title']
            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"
            markdown += f"<div style='color: #34495e; margin-bottom: 20px;'>\n\n{section['Content']}\n\n</div>\n\n"
            
            for subsection in section.get('Subsections', []):
                subsection_number = subsection['SectionNumber']
                subsection_title = subsection['Title']
                markdown += f"<h3 id='{cls.slugify(subsection_title)}' style='color: #34495e;'>{subsection_number} {subsection_title}</h3>\n\n"
                markdown += f"<div style='color: #34495e; margin-bottom: 20px;'>\n\n{subsection['Content']}\n\n</div>\n\n"
        
        markdown += "</div>"
        return markdown
        

router = APIRouter()

class DocumentRequest(BaseModel):
    query: str

class DocumentResponse(BaseModel):
    json_document: Dict
    markdown_document: str

@cache(expire=600*24*7)
@router.post("/generate-document", response_model=DocumentResponse)
async def generate_document_endpoint(request: DocumentRequest):
    ai_client = AIClient()
    document_generator = DocumentGenerator(ai_client)
    
    try:
        # Generate the document
        json_document = document_generator.generate_document(request.query)
        
        # Convert to Markdown
        markdown_document = MarkdownConverter.convert_to_markdown(json_document["Document"])
        
        return DocumentResponse(
            json_document=json_document,
            markdown_document=markdown_document
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

class CacheTestResponse(BaseModel):
    result: str
    execution_time: float

@router.get("/test-cache/{test_id}", response_model=CacheTestResponse)
@cache(expire=60)  # Cache for 1 minute
async def test_cache(test_id: int):
    start_time = time.time()
    
    # Simulate some time-consuming operation
    await asyncio.sleep(2)
    
    result = f"Test result for ID: {test_id}"
    
    end_time = time.time()
    execution_time = end_time - start_time
    
    return CacheTestResponse(
        result=result,
        execution_time=execution_time
    )