Add main.py

main.py

@@ -1,15 +1,13 @@
 import os
 import asyncio
-from typing import TypedDict, Dict
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
-
+from langchain_chroma import Chroma
-from langchain_openai import ChatOpenAI
 from langchain_core.messages import HumanMessage
 from langchain.tools import tool
-from langgraph.graph import StateGraph, START, END
+from deepagents import create_deep_agent
-from pydantic import BaseModel, Field
+from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
-from langchain_core.output_parsers import PydanticOutputParser
 
-# LLM setup
+# LLM configuration – always OpenRouter
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -17,138 +15,47 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
-# State definition
+# Embeddings and vector store for RAG
-class CodeReviewState(TypedDict):
+embeddings = OpenAIEmbeddings(
-    code: str
+    model="text-embedding-3-small",
-    draft_review: str
+    base_url="https://openrouter.ai/api/v1",
-    criteria_scores: Dict[str, int]
+    api_key=os.getenv("OPENAI_API_KEY"),
-    weakest_criterion: str
+)
-    verdict: str
+vector_store = Chroma(collection_name="knowledge", embedding_function=embeddings)
-    round: int
-    max_rounds: int
 
-# Reflect output model
+# Backend for file operations – virtual mode so no real files are created
-class ReflectOutput(BaseModel):
-    pep8: int = Field(description="Score 0-10 for PEP8 compliance")
-    type_hints: int = Field(description="Score 0-10 for type hints usage")
-    edge_cases: int = Field(description="Score 0-10 for edge case coverage")
-    naming: int = Field(description="Score 0-10 for naming conventions")
-    weakest_criterion: str = Field(description="Criterion with lowest score")
-    verdict: str = Field(description="'ok' or 'needs_revision'")
-
-reflect_parser = PydanticOutputParser(pydantic_object=ReflectOutput)
-
-# Dummy tool for agent
-@tool
-def echo_tool(query: str) -> str:
-    return query
-
-# Agent creation
-from deepagents import create_deep_agent
-from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
 backend = CompositeBackend(
     default=LocalShellBackend(root_dir="./workspace", virtual_mode=True, inherit_env=True),
     routes={},
 )
+
+# Tool that performs a vector search and returns the top 3 snippets
+@tool
+def rag_search(query: str) -> str:
+    """Search the vector store for relevant documents and return a concise summary."""
+    docs = vector_store.similarity_search(query, k=3)
+    if not docs:
+        return "No relevant information found."
+    snippets = "\n---\n".join(doc.page_content for doc in docs)
+    return f"Top matches:\n{snippets}"
+
+# Create the deep agent with the RAG tool
 agent = create_deep_agent(
     model=llm,
-    tools=[echo_tool],
+    tools=[rag_search],
     backend=backend,
-    system_prompt="You are a code review assistant.",
+    system_prompt="You are a helpful assistant that uses a knowledge base to answer questions. Use the rag_search tool when you need external information.",
 )
 
-# Node functions
+async def main():
-async def draft_review(state: CodeReviewState) -> CodeReviewState:
+    # Example user query
-    prompt = f"""Write a concise code review (3-6 points) for the following Python function. Focus on style, correctness, and potential improvements.
+    user_query = "What are the main causes of climate change?"
-
+    result = await agent.ainvoke(
-```python
+        {"messages": [HumanMessage(content=user_query)]},
-{state['code']}
+        {"configurable": {"thread_id": "session-1"}},
-```
+    )
-
+    # Print the final assistant message
-Return only the review text."""
+    print(result["messages"][-1].content)
-    response = await agent.ainvoke({"messages": [HumanMessage(content=prompt)]}, {"configurable": {"thread_id": "draft-review"}})
-    review_text = response["messages"][-1].content
-    state["draft_review"] = review_text
-    return state
-
-async def reflect(state: CodeReviewState) -> CodeReviewState:
-    prompt = f"""You are a senior reviewer. Evaluate the following review text against four criteria: PEP8, type hints, edge cases, naming. Assign each a score 0-10. Identify the weakest criterion and give a verdict: 'ok' if all scores >=7, else 'needs_revision'. Return a JSON with keys: pep8, type_hints, edge_cases, naming, weakest_criterion, verdict.
-
-Review:
-{state['draft_review']}"""
-    response = await agent.ainvoke({"messages": [HumanMessage(content=prompt)]}, {"configurable": {"thread_id": "reflect"}})
-    json_text = response["messages"][-1].content
-    try:
-        parsed = reflect_parser.parse(json_text)
-    except Exception:
-        parsed = ReflectOutput(pep8=5, type_hints=5, edge_cases=5, naming=5, weakest_criterion="pep8", verdict="needs_revision")
-    state["criteria_scores"] = {
-        "pep8": parsed.pep8,
-        "type_hints": parsed.type_hints,
-        "edge_cases": parsed.edge_cases,
-        "naming": parsed.naming,
-    }
-    state["weakest_criterion"] = parsed.weakest_criterion
-    state["verdict"] = parsed.verdict
-    return state
-
-async def rewrite(state: CodeReviewState) -> CodeReviewState:
-    crit = state["weakest_criterion"]
-    prompt = f"""Improve the review section that addresses the weakest criterion '{crit}'. Provide a more detailed point for that criterion. Keep the rest of the review unchanged.
-
-Current review:
-{state['draft_review']}"""
-    response = await agent.ainvoke({"messages": [HumanMessage(content=prompt)]}, {"configurable": {"thread_id": "rewrite"}})
-    new_review = response["messages"][-1].content
-    state["draft_review"] = new_review
-    state["round"] += 1
-    return state
-
-# Graph definition
-from langgraph.graph import StateGraph
-
-graph = StateGraph(CodeReviewState)
-
-graph.add_node("draft_review", draft_review)
-graph.add_node("reflect", reflect)
-graph.add_node("rewrite", rewrite)
-
-graph.set_entry_point("draft_review")
-graph.add_edge("draft_review", "reflect")
-
-graph.add_conditional_edges(
-    "reflect",
-    lambda x: "END" if x["verdict"] == "ok" else "rewrite" if x["round"] < x["max_rounds"] else "END",
-)
-
-graph.add_edge("rewrite", "reflect")
-app = graph.compile()
-
-# Demo function
-async def demo():
-    sample_code = """def sort_numbers(arr):
-    return sorted(arr)"""
-    init_state: CodeReviewState = {
-        "code": sample_code,
-        "draft_review": "",
-        "criteria_scores": {},
-        "weakest_criterion": "",
-        "verdict": "",
-        "round": 0,
-        "max_rounds": 2,
-    }
-    result = await app.ainvoke(init_state)
-    print("--- Draft Review ---")
-    print(result["draft_review"])
-    print("\n--- Scores ---")
-    print(result["criteria_scores"])
-    print("\n--- Verdict ---")
-    print(result["verdict"])
-    if result["verdict"] == "needs_revision":
-        print("\n--- Final Review After Rewrite ---")
-        print(result["draft_review"])
-        print("\n--- Final Scores ---")
-        print(result["criteria_scores"])
 
 if __name__ == "__main__":
-    asyncio.run(demo())
+    asyncio.run(main())