Upload models.py

models.py

@@ -0,0 +1,232 @@
+from openai import OpenAI
+import json
+import time
+from datetime import datetime
+from utils import log_step, SYSTEM_PROMPT, logger
+
+class LSP:
+    def __init__(self, api_key: str):
+        if not api_key or api_key.strip() == "":
+            raise ValueError("API key cannot be empty")
+        self.client = OpenAI(api_key=api_key)
+
+    def get_section_map(self, content: str) -> tuple[dict, float]:
+        """Get map of sections without loading full content"""
+        start_time = time.time()
+        log_step("SECTION_MAP", "Starting section map generation")
+        
+        sections = {}
+        in_code_block = False
+        current_section = None
+        section_start = 0
+        
+        lines = content.split('\n')
+        for i, line in enumerate(lines):
+            if line.strip().startswith('```'):
+                in_code_block = not in_code_block
+                continue
+                
+            if not in_code_block and line.strip().startswith('#'):
+                if current_section:
+                    sections[current_section] = (section_start, i-1)
+                current_section = line.strip()
+                section_start = i
+                
+        if current_section:
+            sections[current_section] = (section_start, len(lines)-1)
+        
+        elapsed = time.time() - start_time
+        log_step("SECTION_MAP", f"Section map generated in {elapsed:.2f}s", sections)
+        return sections, elapsed
+
+    def extract_section(self, content: str, start: int, end: int) -> tuple[str, float]:
+        """Get just the content of target section"""
+        start_time = time.time()
+        log_step("EXTRACT", f"Extracting section from line {start} to {end}")
+        
+        lines = content.split('\n')
+        section = '\n'.join(lines[start:end+1])
+        
+        elapsed = time.time() - start_time
+        log_step("EXTRACT", f"Section extracted in {elapsed:.2f}s", {"length": len(section)})
+        return section, elapsed
+
+    def replace_section(self, content: str, start: int, end: int, new_content: str) -> tuple[str, float]:
+        """Replace section and return full content"""
+        start_time = time.time()
+        log_step("REPLACE", f"Replacing section from line {start} to {end}")
+        
+        lines = content.split('\n')
+        lines[start:end+1] = new_content.split('\n')
+        result = '\n'.join(lines)
+        
+        elapsed = time.time() - start_time
+        log_step("REPLACE", f"Section replaced in {elapsed:.2f}s", {"new_length": len(result)})
+        return result, elapsed
+
+    def tools(self):
+        return [
+            {
+                "type": "function",
+                "function": {
+                    "name": "identify_section",
+                    "description": "Identify which section needs to be modified based on headers map",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "section_name": {
+                                "type": "string",
+                                "description": "Name of section to edit"
+                            },
+                            "reason": {
+                                "type": "string",
+                                "description": "Why this section was chosen"
+                            }
+                        },
+                        "required": ["section_name", "reason"]
+                    }
+                }
+            },
+            {
+                "type": "function",
+                "function": {
+                    "name": "modify_section",
+                    "description": "Make changes to the section content",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "modified_content": {
+                                "type": "string",
+                                "description": "New content for the section"
+                            },
+                            "changes_made": {
+                                "type": "string",
+                                "description": "Description of changes"
+                            }
+                        },
+                        "required": ["modified_content", "changes_made"]
+                    }
+                }
+            }
+        ]
+
+    async def edit_smart(self, content: str, instruction: str) -> tuple[str, list[str], float]:
+        """Smart editing with section targeting"""
+        total_start = time.time()
+        timings = {}
+        traces = []
+        
+        try:
+            # Phase 1: Get section map
+            section_map, map_time = self.get_section_map(content)
+            timings['section_map'] = map_time
+            traces.append(f"[{map_time:.2f}s] Found sections: {json.dumps(section_map, indent=2)}")
+
+            # Phase 2: Identify section
+            messages = [
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": f"Section map: {json.dumps(section_map)}\nInstruction: {instruction}"}
+            ]
+
+            identify_start = time.time()
+            response = self.client.chat.completions.create(
+                model="gpt-4o-mini",
+                messages=messages,
+                tools=self.tools()
+            )
+            identify_time = time.time() - identify_start
+            timings['identify_section'] = identify_time
+            
+            tool_call = response.choices[0].message.tool_calls[0]
+            args = json.loads(tool_call.function.arguments)
+            target_section = args["section_name"]
+            start, end = section_map[target_section]
+            
+            traces.append(f"[{identify_time:.2f}s] Selected section: {target_section} (lines {start}-{end})")
+            
+            # Phase 3: Extract section
+            section_content, extract_time = self.extract_section(content, start, end)
+            timings['extract_section'] = extract_time
+            traces.append(f"[{extract_time:.2f}s] Extracted content:\n{section_content}")
+            
+            # Phase 4: Modify section
+            modify_start = time.time()
+            messages.extend([
+                response.choices[0].message,
+                {
+                    "role": "tool",
+                    "content": json.dumps({"success": True, "section": target_section}),
+                    "tool_call_id": tool_call.id
+                },
+                {
+                    "role": "user",
+                    "content": f"Here's the section to modify:\n{section_content}"
+                }
+            ])
+
+            response = self.client.chat.completions.create(
+                model="gpt-4o-mini",
+                messages=messages,
+                tools=self.tools()
+            )
+            modify_time = time.time() - modify_start
+            timings['modify_section'] = modify_time
+
+            tool_call = response.choices[0].message.tool_calls[0]
+            args = json.loads(tool_call.function.arguments)
+            traces.append(f"[{modify_time:.2f}s] Modified content:\n{args['modified_content']}")
+
+            # Phase 5: Replace section
+            result, replace_time = self.replace_section(content, start, end, args["modified_content"])
+            timings['replace_section'] = replace_time
+            
+            total_time = time.time() - total_start
+            timings['total'] = total_time
+            
+            # Add timing summary
+            timing_summary = "\nTiming Summary:\n" + "\n".join([
+                f"- {step}: {time:.2f}s" for step, time in timings.items()
+            ])
+            traces.append(timing_summary)
+            
+            return result, traces, total_time
+
+        except Exception as e:
+            total_time = time.time() - total_start
+            error_msg = str(e)
+            traces.append(f"Error after {total_time:.2f}s: {error_msg}")
+            return error_msg, traces, total_time
+
+    async def edit_naive(self, content: str, instruction: str) -> tuple[str, list[str], float]:
+        """Naive approach - give everything to AI"""
+        total_start = time.time()
+        traces = [f"[0.00s] Starting naive edit (sending entire document)"]
+        
+        try:
+            messages = [
+                {"role": "system", "content": "You are a document editor. Edit the provided document according to instructions."},
+                {"role": "user", "content": f"Edit this document according to this instruction: {instruction}\n\nDocument:\n{content}"}
+            ]
+
+            api_start = time.time()
+            response = self.client.chat.completions.create(
+                model="gpt-4o-mini",
+                messages=messages
+            )
+            api_time = time.time() - api_start
+
+            result = response.choices[0].message.content
+            total_time = time.time() - total_start
+            
+            traces.extend([
+                f"[{api_time:.2f}s] OpenAI API call completed",
+                f"[{total_time:.2f}s] Total processing completed"
+            ])
+            
+            return result, traces, total_time
+            
+        except Exception as e:
+            total_time = time.time() - total_start
+            error_msg = str(e)
+            traces.append(f"Error after {total_time:.2f}s: {error_msg}")
+            return error_msg, traces, total_time