Spaces:

mgbam
/

builder

Sleeping

App Files Files Community

mgbam commited on Jul 18

Commit

c118cfa

verified ·

1 Parent(s): 1280ebf

Update core.py

Browse files

Files changed (1) hide show

core.py +22 -92

core.py CHANGED Viewed

@@ -1,119 +1,49 @@
-# /core.py (Corrected)
-"""
-Core business logic for the code generation application.
-This module orchestrates the entire process from receiving a user query to
-generating the final code. It interacts with the services, extractors, and
-utility modules to fulfill the request.
-"""
 from typing import Dict, List, Optional, Tuple, Generator, Any
-from config import (
-    HTML_SYSTEM_PROMPT, GENERIC_SYSTEM_PROMPT,
-    HTML_SYSTEM_PROMPT_WITH_SEARCH, GENERIC_SYSTEM_PROMPT_WITH_SEARCH,
-    FOLLOW_UP_SYSTEM_PROMPT
-)
 from services import llm_service, search_service
 from extractor import extract_text_from_file, extract_website_content
-from utils import (
-    history_to_messages,
-    remove_code_block,
-    process_image_for_model,
-    apply_search_replace_changes  # <--- FIX: Corrected the function name here
-)
-# --- Type Definitions ---
 History = List[Tuple[Optional[str], Optional[str]]]
 def _determine_system_prompt(language: str, enable_search: bool, history: History) -> Tuple[str, bool]:
-    """Determines the appropriate system prompt based on context."""
-    is_follow_up = False
-    if history and history[-1][1] and ("<!DOCTYPE html>" in history[-1][1] or "<html" in history[-1][1]):
-        is_follow_up = True
-        return FOLLOW_UP_SYSTEM_PROMPT, is_follow_up
     if language == "html":
-        prompt = HTML_SYSTEM_PROMPT_WITH_SEARCH if enable_search else HTML_SYSTEM_PROMPT
     else:
-        base_prompt = GENERIC_SYSTEM_PROMPT_WITH_SEARCH if enable_search else GENERIC_SYSTEM_PROMPT
-        prompt = base_prompt.format(language=language)
-    return prompt, is_follow_up
-def _prepare_user_content(
-    query: str, image_data: Optional[Any], file_path: Optional[str],
-    website_url: Optional[str], enable_search: bool
-) -> any:
-    """Constructs the final user prompt including context from files, web, and search."""
     context_parts = [query]
-    if file_path:
-        file_text = extract_text_from_file(file_path)
-        context_parts.append(f"\n\n--- Reference File Content ---\n{file_text[:8000]}")
-    if website_url:
-        web_text = extract_website_content(website_url)
-        context_parts.append(f"\n\n--- Website Content for Redesign ---\n{web_text[:8000]}")
     full_query = "".join(context_parts)
-    if enable_search and search_service.is_available():
-        search_results = search_service.search(full_query)
-        full_query += f"\n\n--- Web Search Results ---\n{search_results}"
     if image_data is not None:
-        return [
-            {"type": "text", "text": full_query},
-            {"type": "image_url", "image_url": {"url": process_image_for_model(image_data)}}
-        ]
     return full_query
-def generate_code(
-    query: str,
-    image_data: Optional[Any],
-    file_path: Optional[str],
-    website_url: Optional[str],
-    history: History,
-    model_config: Dict[str, str],
-    enable_search: bool,
-    language: str
-) -> Generator[Dict[str, Any], None, None]:
-    """
-    Main generator function to handle a user request and stream responses.
-    """
     system_prompt, is_follow_up = _determine_system_prompt(language, enable_search, history)
     messages = history_to_messages(history, system_prompt)
     user_content = _prepare_user_content(query, image_data, file_path, website_url, enable_search)
     messages.append({'role': 'user', 'content': user_content})
     content_stream = ""
-    stream = llm_service.generate_code_stream(model_config['id'], messages)
-    for chunk in stream:
         content_stream += chunk
-        processed_code = ""
-        if is_follow_up:
-            last_html = history[-1][1] if history and history[-1][1] else ""
-            # <--- FIX: Use the corrected function name here
-            modified_html = apply_search_replace_changes(last_html, content_stream)
-            processed_code = modified_html
-        else:
-            processed_code = remove_code_block(content_stream)
         yield {"code_output": processed_code}
-    # Final processing after stream ends
-    final_content = content_stream
-    if is_follow_up:
-        last_html = history[-1][1] if history and history[-1][1] else ""
-        # <--- FIX: And use the corrected function name here as well
-        final_code = apply_search_replace_changes(last_html, final_content)
-    else:
-        final_code = remove_code_block(final_content)
-    # Use the original query for history display, not the enhanced one.
     history.append((query, final_code))
     yield {"code_output": final_code, "history": history}

+# /core.py
+""" Core business logic for the code generation application. """
 from typing import Dict, List, Optional, Tuple, Generator, Any
+from config import (HTML_SYSTEM_PROMPT, GENERIC_SYSTEM_PROMPT, HTML_SYSTEM_PROMPT_WITH_SEARCH,
+                    GENERIC_SYSTEM_PROMPT_WITH_SEARCH, FollowUpSystemPrompt)
 from services import llm_service, search_service
 from extractor import extract_text_from_file, extract_website_content
+from utils import (history_to_messages, remove_code_block, process_image_for_model,
+                   apply_search_replace_changes)
 History = List[Tuple[Optional[str], Optional[str]]]
 def _determine_system_prompt(language: str, enable_search: bool, history: History) -> Tuple[str, bool]:
+    is_follow_up = bool(history and history[-1][1] and ("<html" in history[-1][1]))
+    if is_follow_up: return FollowUpSystemPrompt, True
     if language == "html":
+        return (HTML_SYSTEM_PROMPT_WITH_SEARCH if enable_search else HTML_SYSTEM_PROMPT), False
     else:
+        base = GENERIC_SYSTEM_PROMPT_WITH_SEARCH if enable_search else GENERIC_SYSTEM_PROMPT
+        return base.format(language=language), False
+def _prepare_user_content(query: str, image_data: Optional[Any], file_path: Optional[str], website_url: Optional[str], enable_search: bool) -> Any:
     context_parts = [query]
+    if file_path: context_parts.append(f"\n\n--- Reference File ---\n{extract_text_from_file(file_path)[:8000]}")
+    if website_url: context_parts.append(f"\n\n--- Website to Redesign ---\n{extract_website_content(website_url)}")
     full_query = "".join(context_parts)
+    if enable_search and search_service.is_available(): full_query += f"\n\n--- Web Search Results ---\n{search_service.search(full_query)}"
     if image_data is not None:
+        return [{"type": "text", "text": full_query}, {"type": "image_url", "image_url": {"url": process_image_for_model(image_data)}}]
     return full_query
+def generate_code(query: str, image_data: Optional[Any], file_path: Optional[str], website_url: Optional[str], history: History,
+                  model_config: Dict[str, str], enable_search: bool, language: str) -> Generator[Dict[str, Any], None, None]:
     system_prompt, is_follow_up = _determine_system_prompt(language, enable_search, history)
     messages = history_to_messages(history, system_prompt)
     user_content = _prepare_user_content(query, image_data, file_path, website_url, enable_search)
     messages.append({'role': 'user', 'content': user_content})
     content_stream = ""
+    for chunk in llm_service.generate_code_stream(model_config['id'], messages):
         content_stream += chunk
+        processed_code = apply_search_replace_changes(history[-1][1], content_stream) if is_follow_up else remove_code_block(content_stream)
         yield {"code_output": processed_code}
+    final_code = apply_search_replace_changes(history[-1][1], content_stream) if is_follow_up else remove_code_block(content_stream)
     history.append((query, final_code))
     yield {"code_output": final_code, "history": history}