Orchestrating Multi‑Agent Systems with Long‑Term Memory for Complex Autonomous Software‑Engineering Workflows

# orchestrator.py
import os
import uuid
import json
from datetime import datetime
from typing import List, Dict

from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
from pinecone import Pinecone
from neo4j import GraphDatabase

# -------------------- Configuration --------------------
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
NEO4J_URI = os.getenv("NEO4J_URI")
NEO4J_USER = os.getenv("NEO4J_USER")
NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD")

# -------------------- LLM Wrapper --------------------
llm = OpenAI(api_key=OPENAI_API_KEY, temperature=0.2)

# -------------------- Vector Store --------------------
pc = Pinecone(api_key=PINECONE_API_KEY)
index = pc.Index("se-knowledge")
embedding_dim = 1536  # OpenAI text-embedding-ada-002

def embed(text: str) -> List[float]:
    """Return a 1536‑dimensional embedding for the given text."""
    resp = llm.embeddings.create(input=[text], model="text-embedding-ada-002")
    return resp['data'][0]['embedding']

def upsert_memory(artifact_id: str, text: str, metadata: Dict):
    vec = embed(text)
    index.upsert(vectors=[{
        "id": artifact_id,
        "values": vec,
        "metadata": metadata
    }])

def query_memory(query: str, top_k: int = 5) -> List[Dict]:
    q_vec = embed(query)
    results = index.query(vector=q_vec, top_k=top_k, include_metadata=True)
    return results['matches']

# -------------------- Knowledge Graph --------------------
driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))

def create_artifact_node(tx, artifact_id, artifact_type, name):
    tx.run(
        """
        MERGE (a:Artifact {id: $artifact_id})
        SET a.type = $artifact_type,
            a.name = $name,
            a.created = datetime()
        """,
        artifact_id=artifact_id,
        artifact_type=artifact_type,
        name=name,
    )

def create_relationship(tx, src_id, dst_id, rel_type):
    tx.run(
        """
        MATCH (src:Artifact {id: $src_id}), (dst:Artifact {id: $dst_id})
        MERGE (src)-[r:%s]->(dst)
        SET r.since = datetime()
        """ % rel_type,
        src_id=src_id,
        dst_id=dst_id,
    )

def graph_add_artifact(artifact_id, artifact_type, name):
    with driver.session() as session:
        session.write_transaction(create_artifact_node, artifact_id, artifact_type, name)

def graph_add_relation(src_id, dst_id, rel_type):
    with driver.session() as session:
        session.write_transaction(create_relationship, src_id, dst_id, rel_type)

# -------------------- Agent Implementations --------------------
def feature_analyst(user_story: str) -> Dict:
    """Extract functional specs from a natural‑language story."""
    prompt = PromptTemplate(
        template="You are a software analyst. Convert the following user story into a JSON spec with fields: `summary`, `acceptance_criteria`, and `constraints`.\nUser Story: {story}",
        input_variables=["story"]
    )
    response = llm(prompt.format(story=user_story))
    spec = json.loads(response)
    uid = str(uuid.uuid4())
    upsert_memory(uid, json.dumps(spec), {"type": "feature_spec"})
    graph_add_artifact(uid, "FeatureSpec", spec["summary"])
    return {"id": uid, "spec": spec}

def architecture_synthesizer(feature_spec: Dict) -> Dict:
    """Propose a simple RESTful API contract."""
    prompt = PromptTemplate(
        template="""
        You are an architect. Based on the feature spec below, generate a minimal OpenAPI v3 contract.
        Spec: {spec}
        Return ONLY the JSON representation.
        """,
        input_variables=["spec"]
    )
    response = llm(prompt.format(spec=json.dumps(feature_spec["spec"])))
    contract = json.loads(response)
    uid = str(uuid.uuid4())
    upsert_memory(uid, json.dumps(contract), {"type": "api_contract"})
    graph_add_artifact(uid, "APIContract", f"Contract for {feature_spec['spec']['summary']}")
    graph_add_relation(feature_spec["id"], uid, "DESCRIBES")
    return {"id": uid, "contract": contract}

def code_generator(contract: Dict) -> List[Dict]:
    """Generate backend (FastAPI) and frontend (React) stubs."""
    generated = []
    for lang, template in [("python", "FastAPI stub"), ("typescript", "React component")]:
        prompt = PromptTemplate(
            template="""
            Generate a {lang} implementation that satisfies the given OpenAPI contract.
            Contract: {contract}
            Return ONLY the source code as a fenced block.
            """,
            input_variables=["lang", "contract"]
        )
        response = llm(prompt.format(lang=lang, contract=json.dumps(contract["contract"])))
        code = response.strip("```")
        uid = str(uuid.uuid4())
        upsert_memory(uid, code, {"type": "source_code", "language": lang})
        graph_add_artifact(uid, "SourceCode", f"{lang} stub")
        graph_add_relation(contract["id"], uid, "GENERATES")
        generated.append({"id": uid, "language": lang, "code": code})
    return generated

def test_engineer(source_artifacts: List[Dict]) -> List[Dict]:
    """Create pytest and jest test suites."""
    tests = []
    for src in source_artifacts:
        prompt = PromptTemplate(
            template="""
            Write a test suite for the following {lang} code.
            Code:
            {code}
            Use pytest for Python and jest for TypeScript. Return only the test file contents.
            """,
            input_variables=["lang", "code"]
        )
        response = llm(prompt.format(lang=src["language"], code=src["code"]))
        test_code = response.strip("```")
        uid = str(uuid.uuid4())
        upsert_memory(uid, test_code, {"type": "test_suite", "language": src["language"]})
        graph_add_artifact(uid, "TestSuite", f"{src['language']} tests")
        graph_add_relation(src["id"], uid, "VALIDATES")
        tests.append({"id": uid, "language": src["language"], "code": test_code})
    return tests

def ci_orchestrator(test_artifacts: List[Dict]) -> bool:
    """Run a mock CI step – in reality you would invoke Docker or GitHub Actions."""
    # For brevity we simulate pass/fail based on a random heuristic.
    import random
    success = all(random.random() > 0.1 for _ in test_artifacts)  # 90% chance of pass
    # Store CI log in LTM
    log = f"CI run at {datetime.utcnow().isoformat()} – {'PASS' if success else 'FAIL'}"
    uid = str(uuid.uuid4())
    upsert_memory(uid, log, {"type": "ci_log"})
    graph_add_artifact(uid, "CILog", "CI outcome")
    for test in test_artifacts:
        graph_add_relation(test["id"], uid, "TRIGGERED")
    return success

def release_manager(ci_success: bool) -> None:
    if not ci_success:
        raise RuntimeError("CI failed – aborting deployment")
    # Simulate deployment
    deployment_id = str(uuid.uuid4())
    log = f"Deployed to staging at {datetime.utcnow().isoformat()}"
    upsert_memory(deployment_id, log, {"type": "deployment"})
    graph_add_artifact(deployment_id, "Deployment", "Staging release")
    print("[Release] ", log)

# -------------------- Top‑Level Planner --------------------
def orchestrate_feature(user_story: str):
    # 1. Analyze feature
    feature = feature_analyst(user_story)

    # 2. Synthesize architecture
    arch = architecture_synthesizer(feature)

    # 3. Generate code
    code_artifacts = code_generator(arch)

    # 4. Produce tests
    test_artifacts = test_engineer(code_artifacts)

    # 5. Run CI
    ci_ok = ci_orchestrator(test_artifacts)

    # 6. Deploy if CI passes
    release_manager(ci_ok)

if __name__ == "__main__":
    sample_story = """
    As a user, I want to be able to register an account using email and password,
    receive a verification email, and log in securely.
    """
    orchestrate_feature(sample_story)

class MemoryBackedAgent:
    def __init__(self, name: str, llm, vector_index, graph_driver):
        self.name = name
        self.llm = llm
        self.index = vector_index
        self.graph = graph_driver

    def remember(self, text: str, metadata: dict) -> str:
        """Persist a piece of knowledge and return its UUID."""
        uid = str(uuid.uuid4())
        vec = embed(text)                     # reuse the embed() helper
        self.index.upsert(vectors=[{
            "id": uid,
            "values": vec,
            "metadata": metadata
        }])
        with self.graph.session() as s:
            s.write_transaction(
                create_artifact_node, uid, metadata.get("type", "generic"), self.name
            )
        return uid

    def recall(self, query: str, top_k: int = 3) -> List[dict]:
        """Semantic retrieval of relevant memories."""
        results = self.index.query(vector=embed(query), top_k=top_k, include_metadata=True)
        return results["matches"]

def bulk_upsert(artifacts: List[Tuple[str, str, dict]]):
    """
    artifacts: List of (uuid, text, metadata) tuples.
    """
    vectors = [
        {"id": uid, "values": embed(txt), "metadata": meta}
        for uid, txt, meta in artifacts
    ]
    index.upsert(vectors=vectors)

def planner_prompt(goal: str, agents: List[str]) -> List[dict]:
    template = """
    You are a planner for an autonomous software‑engineering team.
    Goal: {goal}
    Available agents: {agents}
    Produce a JSON array where each element has:
      - "agent": name of the agent to invoke
      - "input": the text you will feed to that agent
    Keep the order logical and include any necessary intermediate steps.
    """
    prompt = PromptTemplate(template=template, input_variables=["goal", "agents"])
    raw = llm(prompt.format(goal=goal, agents=", ".join(agents)))
    return json.loads(raw)

# Example usage:
schedule = planner_prompt(
    goal="Implement password‑reset flow with email verification",
    agents=["Feature Analyst", "Architecture Synthesizer", "Code Generator", "Test Engineer", "CI Orchestrator"]
)