app/services/context_formatter.py

from app.models.graph_context import GraphContext


class ContextFormatter:

    def format(
        self,
        context: GraphContext,
    ) -> str:

        seed_names = {
            node.qualified_name
            for node in context.nodes
            if node.properties.get("source_code")
        }

        lines = []

        lines.append("# Relevant Code")
        lines.append("")

        for node in context.nodes:

            source = node.properties.get(
                "source_code"
            )

            if not source:
                continue

            lines.append(
                f"## {node.qualified_name}"
            )

            lines.append("")

            lines.append("```python")
            lines.append(source)
            lines.append("```")
            lines.append("")

        lines.append("# Related Graph Nodes")
        lines.append("")

        for node in context.nodes:

            if node.qualified_name in seed_names:
                continue

            props = []

            if node.properties.get("name"):

                props.append(
                    f"name={node.properties['name']}"
                )

            if node.properties.get("module"):

                props.append(
                    f"module={node.properties['module']}"
                )

            label = ",".join(node.labels)

            lines.append(
                f"- {node.qualified_name} [{label}] ({', '.join(props)})"
            )

        lines.append("")
        lines.append("# Relationships")
        lines.append("")

        for edge in context.edges:

            lines.append(
                f"- {edge.source} --{edge.type}--> {edge.target}"
            )

        return "\n".join(lines)