Merge pull request #132 from equinor/create-decision-tree

Create decision tree

Author: jafr67

poetry.lock

@@ -33,6 +33,7 @@ asgi-lifespan = "^2.1.0"
 pyagrum = "^2.2.1"
 black = "^25.9.0"
 flake8 = "^7.3.0"
+networkx = "^3.5"
 
 [tool.pytest.ini_options]
 asyncio_mode = "strict"

pyproject.toml

@@ -19,6 +19,7 @@
 from src.services.option_service import OptionService
 from src.services.user_service import UserService
 from src.services.solver_service import SolverService
+from src.services.structure_service import StructureService
 from src.config import config
 from src.database import get_connection_string_and_token, build_connection_url
 
@@ -110,3 +111,7 @@ async def get_user_service() -> UserService:
 
 async def get_solver_service() -> SolverService:
     return SolverService(await get_scenario_service())
+
+
+async def get_structure_service() -> StructureService:
+    return StructureService(await get_scenario_service())

src/dependencies.py

@@ -0,0 +1,20 @@
+import uuid
+from typing import List, Optional
+from pydantic import BaseModel, Field
+from src.dtos.issue_dtos import IssueOutgoingDto
+
+
+class EdgeUUIDDto(BaseModel):
+    tail: uuid.UUID
+    head: uuid.UUID | None
+
+
+class EndPointNodeDto(BaseModel):
+    id: uuid.UUID = Field(default_factory=uuid.uuid4)
+    scenario_id: uuid.UUID
+    type: str = "EndPoint"
+
+
+class DecisionTreeDTO(BaseModel):
+    tree_node: IssueOutgoingDto | EndPointNodeDto
+    children: Optional[List["DecisionTreeDTO"]] = None

src/dtos/decision_tree_dtos.py

@@ -18,6 +18,8 @@ class EdgeIncomingDto(EdgeDto):
 class EdgeOutgoingDto(EdgeDto):
     head_node: NodeOutgoingDto
     tail_node: NodeOutgoingDto
+    head_issue_id: uuid.UUID
+    tail_issue_id: uuid.UUID
 
 
 class EdgeMapper:
@@ -28,6 +30,8 @@ def to_outgoing_dto(entity: Edge) -> EdgeOutgoingDto:
             tail_id=entity.tail_id,
             head_id=entity.head_id,
             scenario_id=entity.scenario_id,
+            head_issue_id=entity.head_node.issue_id,
+            tail_issue_id=entity.tail_node.issue_id,
             head_node=NodeMapper.to_outgoing_dto(entity.head_node),
             tail_node=NodeMapper.to_outgoing_dto(entity.tail_node),
         )

src/dtos/edge_dtos.py

@@ -29,6 +29,7 @@
     NodeOutgoingDto,
     NodeViaIssueOutgoingDto,
 )
+
 from src.constants import DatabaseConstants, DepricatedIssueTypes
 
 

src/dtos/issue_dtos.py

@@ -18,6 +18,7 @@
 import src.routes.outcome_routes as outcome_routes
 import src.routes.option_routes as option_routes
 import src.routes.solver_routes as solver_routes
+import src.routes.structure_routes as structure_routes
 from src.config import config
 from src.session_manager import sessionmanager
 from fastapi.middleware.cors import CORSMiddleware
@@ -71,6 +72,7 @@ async def root():
 app.include_router(edge_routes.router, dependencies=[Depends(verify_token)])
 app.include_router(outcome_routes.router, dependencies=[Depends(verify_token)])
 app.include_router(option_routes.router, dependencies=[Depends(verify_token)])
+app.include_router(structure_routes.router, dependencies=[Depends(verify_token)])
 
 if __name__ == "__main__":
     uvicorn.run("src.main:app", port=8080)

src/main.py

@@ -0,0 +1,19 @@
+import uuid
+from typing import Optional
+from fastapi import APIRouter, Depends, HTTPException
+from src.services.structure_service import StructureService
+from src.dependencies import get_structure_service
+from src.dtos.decision_tree_dtos import DecisionTreeDTO
+
+
+router = APIRouter(tags=["structure"])
+
+@router.get("/structure/{scenario_id}/decision_tree")
+async def get_decision_tree(
+    scenario_id: uuid.UUID,
+    structure_service: StructureService = Depends(get_structure_service)
+) -> Optional[DecisionTreeDTO]:
+    try:
+        return await structure_service.create_decision_tree_dtos(scenario_id)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))