Add application file

Dockerfile

@@ -10,4 +10,4 @@ COPY --chown=user ./requirements.txt requirements.txt
 RUN pip install --no-cache-dir --upgrade -r requirements.txt
 
 COPY --chown=user . /app
-CMD ["gunicorn", "--bind", "0.0.0.0:7860", "app:app"]
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -1,68 +1,153 @@
-from flask import Flask, jsonify, render_template, request, Response, stream_with_context
-from research_generator import ResearchGenerator
+from logging import Logger
+import os
+from datetime import datetime
+from typing import Dict, List, Optional
+from dotenv import load_dotenv
+from langchain_ollama import OllamaLLM
+from langgraph.graph import END, StateGraph
+from models import AgentState, DiagramType
+from config import Config
+from langchain_core.prompts import ChatPromptTemplate
+from file_manager import FileManager
+from agents.analysis import (
+    EntitiesAnalysis,
+    ActorsAnalysis,
+    SequenceAnalysis,
+    StateAnalysis,
+    ArchitectureAnalysis,
+    DeploymentAnalysis
+)
+from agents.code import (
+    ClassDiagramGenerator,
+    UseCaseDiagramGenerator,
+    SequenceDiagramGenerator,
+    ActivityDiagramGenerator,
+    ComponentDiagramGenerator,
+    DeploymentDiagramGenerator,
+    StateDiagramGenerator,
+    TimingDiagramGenerator,
+    ObjectDiagramGenerator
+)
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
 import asyncio
+from langchain_core.callbacks import StreamingStdOutCallbackHandler
 import json
-from langchain_core.output_parsers import StrOutputParser
-from langchain_core.prompts import ChatPromptTemplate
 
-app = Flask(__name__)
-generator = ResearchGenerator()
-
-@app.route('/')
-def index():
-    return {}
-
-async def process_state(state):
-    if "index" in state and state["index"]:
-        # If we have an index, process it with the LLM
-        prompt = ChatPromptTemplate.from_template("""
-        Process the following research state:
-        {state}
-        """)
-        chain = prompt | generator.llm | StrOutputParser()
-        
-        async for chunk in chain.astream({"state": json.dumps(state)}):
-            yield f"data: {json.dumps({'chunk': chunk, 'state': state})}\n\n"
-    else:
-        # If we're generating the index, stream it directly
-        yield f"data: {json.dumps({'state': state})}\n\n"
-
-@app.route('/generate', methods=['POST'])
-def generate_research():
-    data = request.get_json()
-    subject = data.get('subject')
-    if not subject:
-        return jsonify({'error': 'Subject is required'}), 400
-
-    async def generate():
-        generator = ResearchGenerator()
-        async for state in generator.astream({
-            "subject": subject,
-            "index": [],
-            "content": {},
-            "current_step": 0
-        }):
-            async for chunk in process_state(state):
-                yield chunk
-
-    def sync_generate():
-        loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(loop)
-        try:
-            async_gen = generate()
-            while True:
-                try:
-                    chunk = loop.run_until_complete(async_gen.__anext__())
-                    yield chunk
-                except StopAsyncIteration:
-                    break
-        finally:
-            loop.close()
-
-    return Response(
-        stream_with_context(sync_generate()),
-        mimetype='text/event-stream'
-    )
-
-if __name__ == '__main__':
-    app.run(debug=True) 
+# Load environment variables
+load_dotenv()
+
+# Initialize FastAPI app
+app = FastAPI()
+
+# Add CORS middleware configuration
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+# Request model
+class DiagramRequest(BaseModel):
+    project_name: str
+    project_description: str
+    selected_diagrams: List[int]
+    output_dir: Optional[str] = "output"
+
+
+
+
+@app.post("/generate-diagrams")
+async def generate_diagrams(request: DiagramRequest):
+    try:
+        initial_state = AgentState(
+            project_name=request.project_name,
+            project_description=request.project_description,
+            output_dir=request.output_dir,
+            selected_diagrams=request.selected_diagrams,
+            entities_classes="",
+            actors_use_cases="",
+            sequence_interactions="",
+            class_diagram="",
+            use_case_diagram="",
+            sequence_diagram="",
+            activity_diagram="",
+            component_diagram="",
+            deployment_diagram="",
+            state_diagram="",
+            timing_diagram="",
+            object_diagram=""
+        )
+
+    
+        diagram_requirements = {
+            DiagramType.CLASS: ["extract_entities"],
+            DiagramType.USE_CASE: ["extract_actors"],
+            DiagramType.SEQUENCE: ["extract_actors", "extract_sequence"],
+            DiagramType.OBJECT: ["extract_entities"],
+            DiagramType.ACTIVITY: ["extract_actors"],
+            DiagramType.COMPONENT: ["extract_entities", "extract_architecture"],
+            DiagramType.DEPLOYMENT: ["extract_architecture", "extract_deployment"],
+            DiagramType.STATE: ["extract_entities", "extract_sequence", "extract_states"],
+            DiagramType.TIMING: ["extract_sequence"]
+        }
+
+       
+        analysis_functions = {
+            "extract_entities": EntitiesAnalysis(),
+            "extract_actors": ActorsAnalysis(),
+            "extract_sequence": SequenceAnalysis(),
+            "extract_states": StateAnalysis(),
+            "extract_architecture": ArchitectureAnalysis(),
+            "extract_deployment": DeploymentAnalysis()
+        }
+
+        generation_functions = {
+            DiagramType.CLASS: ("generate_class", ClassDiagramGenerator()),
+            DiagramType.USE_CASE: ("generate_use_case", UseCaseDiagramGenerator()),
+            DiagramType.SEQUENCE: ("generate_sequence", SequenceDiagramGenerator()),
+            DiagramType.ACTIVITY: ("generate_activity", ActivityDiagramGenerator()),
+            DiagramType.COMPONENT: ("generate_component", ComponentDiagramGenerator()),
+            DiagramType.DEPLOYMENT: ("generate_deployment", DeploymentDiagramGenerator()),
+            DiagramType.STATE: ("generate_state", StateDiagramGenerator()),
+            DiagramType.TIMING: ("generate_timing", TimingDiagramGenerator()),
+            DiagramType.OBJECT: ("generate_object", ObjectDiagramGenerator())
+        }
+
+      
+        selected_types = [DiagramType(d) for d in request.selected_diagrams]
+        added_analysis_steps = set()
+        steps = []
+
+        for diagram_type in selected_types:
+            required_analysis = diagram_requirements.get(diagram_type, [])
+            for step in required_analysis:
+                if step not in added_analysis_steps:
+                    steps.append((step, analysis_functions[step]))
+                    added_analysis_steps.add(step)
+
+            if diagram_type in generation_functions:
+                steps.append(generation_functions[diagram_type])
+
+       
+        async def event_stream():
+            for step_name, step_func in steps:
+                # Determine if it's an analysis step or generation step
+                step_type = "analysis" if step_name in analysis_functions else "generation"
+                # Determine file extension based on step type
+                file_ext = "pu" if step_type == "generation" else "md"
+                yield f'data: {{"type": "step", "step_type": "{step_type}", "file_name": "{step_name}.{file_ext}"}}\n\n'
+                async for chunk in step_func(initial_state):
+                    yield f"data: {chunk}\n\n"
+
+        return StreamingResponse(event_stream(), media_type="text/event-stream")
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+

main.py

@@ -1,153 +0,0 @@
-from logging import Logger
-import os
-from datetime import datetime
-from typing import Dict, List, Optional
-from dotenv import load_dotenv
-from langchain_ollama import OllamaLLM
-from langgraph.graph import END, StateGraph
-from models import AgentState, DiagramType
-from config import Config
-from langchain_core.prompts import ChatPromptTemplate
-from file_manager import FileManager
-from agents.analysis import (
-    EntitiesAnalysis,
-    ActorsAnalysis,
-    SequenceAnalysis,
-    StateAnalysis,
-    ArchitectureAnalysis,
-    DeploymentAnalysis
-)
-from agents.code import (
-    ClassDiagramGenerator,
-    UseCaseDiagramGenerator,
-    SequenceDiagramGenerator,
-    ActivityDiagramGenerator,
-    ComponentDiagramGenerator,
-    DeploymentDiagramGenerator,
-    StateDiagramGenerator,
-    TimingDiagramGenerator,
-    ObjectDiagramGenerator
-)
-from fastapi import FastAPI, HTTPException
-from fastapi.responses import StreamingResponse
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-import asyncio
-from langchain_core.callbacks import StreamingStdOutCallbackHandler
-import json
-
-# Load environment variables
-load_dotenv()
-
-# Initialize FastAPI app
-app = FastAPI()
-
-# Add CORS middleware configuration
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-
-# Request model
-class DiagramRequest(BaseModel):
-    project_name: str
-    project_description: str
-    selected_diagrams: List[int]
-    output_dir: Optional[str] = "output"
-
-
-
-
-@app.post("/generate-diagrams")
-async def generate_diagrams(request: DiagramRequest):
-    try:
-        initial_state = AgentState(
-            project_name=request.project_name,
-            project_description=request.project_description,
-            output_dir=request.output_dir,
-            selected_diagrams=request.selected_diagrams,
-            entities_classes="",
-            actors_use_cases="",
-            sequence_interactions="",
-            class_diagram="",
-            use_case_diagram="",
-            sequence_diagram="",
-            activity_diagram="",
-            component_diagram="",
-            deployment_diagram="",
-            state_diagram="",
-            timing_diagram="",
-            object_diagram=""
-        )
-
-    
-        diagram_requirements = {
-            DiagramType.CLASS: ["extract_entities"],
-            DiagramType.USE_CASE: ["extract_actors"],
-            DiagramType.SEQUENCE: ["extract_actors", "extract_sequence"],
-            DiagramType.OBJECT: ["extract_entities"],
-            DiagramType.ACTIVITY: ["extract_actors"],
-            DiagramType.COMPONENT: ["extract_entities", "extract_architecture"],
-            DiagramType.DEPLOYMENT: ["extract_architecture", "extract_deployment"],
-            DiagramType.STATE: ["extract_entities", "extract_sequence", "extract_states"],
-            DiagramType.TIMING: ["extract_sequence"]
-        }
-
-       
-        analysis_functions = {
-            "extract_entities": EntitiesAnalysis(),
-            "extract_actors": ActorsAnalysis(),
-            "extract_sequence": SequenceAnalysis(),
-            "extract_states": StateAnalysis(),
-            "extract_architecture": ArchitectureAnalysis(),
-            "extract_deployment": DeploymentAnalysis()
-        }
-
-        generation_functions = {
-            DiagramType.CLASS: ("generate_class", ClassDiagramGenerator()),
-            DiagramType.USE_CASE: ("generate_use_case", UseCaseDiagramGenerator()),
-            DiagramType.SEQUENCE: ("generate_sequence", SequenceDiagramGenerator()),
-            DiagramType.ACTIVITY: ("generate_activity", ActivityDiagramGenerator()),
-            DiagramType.COMPONENT: ("generate_component", ComponentDiagramGenerator()),
-            DiagramType.DEPLOYMENT: ("generate_deployment", DeploymentDiagramGenerator()),
-            DiagramType.STATE: ("generate_state", StateDiagramGenerator()),
-            DiagramType.TIMING: ("generate_timing", TimingDiagramGenerator()),
-            DiagramType.OBJECT: ("generate_object", ObjectDiagramGenerator())
-        }
-
-      
-        selected_types = [DiagramType(d) for d in request.selected_diagrams]
-        added_analysis_steps = set()
-        steps = []
-
-        for diagram_type in selected_types:
-            required_analysis = diagram_requirements.get(diagram_type, [])
-            for step in required_analysis:
-                if step not in added_analysis_steps:
-                    steps.append((step, analysis_functions[step]))
-                    added_analysis_steps.add(step)
-
-            if diagram_type in generation_functions:
-                steps.append(generation_functions[diagram_type])
-
-       
-        async def event_stream():
-            for step_name, step_func in steps:
-                # Determine if it's an analysis step or generation step
-                step_type = "analysis" if step_name in analysis_functions else "generation"
-                # Determine file extension based on step type
-                file_ext = "pu" if step_type == "generation" else "md"
-                yield f'data: {{"type": "step", "step_type": "{step_type}", "file_name": "{step_name}.{file_ext}"}}\n\n'
-                async for chunk in step_func(initial_state):
-                    yield f"data: {chunk}\n\n"
-
-        return StreamingResponse(event_stream(), media_type="text/event-stream")
-
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-

run.sh

@@ -1 +1 @@
-uvicorn main:app --reload
+uvicorn app:app --reload