Spaces:

viraj-421
/

Agentic-AI-Document-Review-System

Paused

Agentic-AI-Document-Review-System / agent_workflow.py

Upload folder using huggingface_hub

f78aa33 verified 4 months ago

4.32 kB

import os
import time
from typing import TypedDict, Literal
from langgraph.graph import StateGraph, END
from pydantic import BaseModel, Field

# Define the structured output for the Critic
class ReviewResponse(BaseModel):
    status: Literal["approved", "rejected"] = Field(
        description="Whether the document meets all criteria or needs revision."
    )
    feedback: str = Field(
        description="Detailed feedback on what is missing or needs improvement. If approved, provide a brief summary of why."
    )

# Define the State
class WorkflowState(TypedDict):
    raw_input: str
    current_draft: str
    review_feedback: str
    status: str
    revision_count: int

# Nodes
def writer_node(state: WorkflowState):
    # Simulate thinking time
    time.sleep(2)
    
    raw_input = state.get("raw_input", "")
    revision_count = state.get("revision_count", 0)

    # Mock PRD generation based on input
    mock_prd = f"""# Product Requirement Document (PRD)

## Executive Summary
This document outlines the requirements based on the raw notes: "{raw_input[:50]}..."

## User Personas
- **Primary User**: General users looking for seamless interaction.
- **Admin**: System administrators managing access.

## User Stories
1. As a user, I want to log in securely so that I can access my dashboard.
2. As an admin, I want to view active sessions.

## Technical Specifications
- **Frontend**: HTML/CSS/JS (Vanilla & Glassmorphism)
- **Backend**: FastAPI with WebSockets
- **Architecture**: Multi-Agent LangGraph

## Edge Cases
- Invalid credentials gracefully handled.
- Rate limits triggered and warnings shown.

## Out of Scope
- Third-party SSO (except where specified).
- Mobile App development.

> **Note**: This is a mock PRD generated locally without an OpenAI API key. Revision count: {revision_count + 1}
"""
    
    return {
        "current_draft": mock_prd,
        "revision_count": revision_count + 1
    }

def critic_node(state: WorkflowState):
    # Simulate thinking time
    time.sleep(2)
    
    revision_count = state.get("revision_count", 0)
    
    # Mock Review Logic: Reject the first time to show the loop, approve the second time.
    if revision_count < 2:
        return {
            "status": "rejected",
            "review_feedback": "The Technical Specifications section needs more detail regarding the database schema. Please revise."
        }
    else:
        return {
            "status": "approved",
            "review_feedback": "The PRD looks excellent and meets all company standards. Moving to human approval."
        }

def human_approval_node(state: WorkflowState):
    # This node doesn't actually do anything automatically.
    # It serves as a placeholder for the graph to pause.
    # When the human approves via API, we will just pass the state through.
    return {"status": "human_approved"}

# Routing
def route_after_critic(state: WorkflowState):
    if state["status"] == "approved":
        return "human_approval"
    elif state["revision_count"] >= 3:
        # Prevent infinite loops, force human review if too many revisions
        return "human_approval"
    else:
        return "writer"

def build_workflow():
    workflow = StateGraph(WorkflowState)
    
    # Add nodes
    workflow.add_node("writer", writer_node)
    workflow.add_node("critic", critic_node)
    workflow.add_node("human_approval", human_approval_node)
    
    # Set entry point
    workflow.set_entry_point("writer")
    
    # Add edges
    workflow.add_edge("writer", "critic")
    workflow.add_conditional_edges(
        "critic",
        route_after_critic,
        {
            "writer": "writer",
            "human_approval": "human_approval"
        }
    )
    workflow.add_edge("human_approval", END)
    
    # Compile graph with a checkpointer for human-in-the-loop (memory)
    from langgraph.checkpoint.memory import MemorySaver
    memory = MemorySaver()
    
    # The graph will interrupt BEFORE the human_approval node
    app = workflow.compile(
        checkpointer=memory,
        interrupt_before=["human_approval"]
    )
    
    return app

# Singleton instance of the graph
workflow_app = build_workflow()

	import os
	import time
	from typing import TypedDict, Literal
	from langgraph.graph import StateGraph, END
	from pydantic import BaseModel, Field

	# Define the structured output for the Critic
	class ReviewResponse(BaseModel):
	status: Literal["approved", "rejected"] = Field(
	description="Whether the document meets all criteria or needs revision."
	)
	feedback: str = Field(
	description="Detailed feedback on what is missing or needs improvement. If approved, provide a brief summary of why."
	)

	# Define the State
	class WorkflowState(TypedDict):
	raw_input: str
	current_draft: str
	review_feedback: str
	status: str
	revision_count: int

	# Nodes
	def writer_node(state: WorkflowState):
	# Simulate thinking time
	time.sleep(2)

	raw_input = state.get("raw_input", "")
	revision_count = state.get("revision_count", 0)

	# Mock PRD generation based on input
	mock_prd = f"""# Product Requirement Document (PRD)

	## Executive Summary
	This document outlines the requirements based on the raw notes: "{raw_input[:50]}..."

	## User Personas
	- Primary User: General users looking for seamless interaction.
	- Admin: System administrators managing access.

	## User Stories
	1. As a user, I want to log in securely so that I can access my dashboard.
	2. As an admin, I want to view active sessions.

	## Technical Specifications
	- Frontend: HTML/CSS/JS (Vanilla & Glassmorphism)
	- Backend: FastAPI with WebSockets
	- Architecture: Multi-Agent LangGraph

	## Edge Cases
	- Invalid credentials gracefully handled.
	- Rate limits triggered and warnings shown.

	## Out of Scope
	- Third-party SSO (except where specified).
	- Mobile App development.

	> Note: This is a mock PRD generated locally without an OpenAI API key. Revision count: {revision_count + 1}
	"""

	return {
	"current_draft": mock_prd,
	"revision_count": revision_count + 1
	}

	def critic_node(state: WorkflowState):
	# Simulate thinking time
	time.sleep(2)

	revision_count = state.get("revision_count", 0)

	# Mock Review Logic: Reject the first time to show the loop, approve the second time.
	if revision_count < 2:
	return {
	"status": "rejected",
	"review_feedback": "The Technical Specifications section needs more detail regarding the database schema. Please revise."
	}
	else:
	return {
	"status": "approved",
	"review_feedback": "The PRD looks excellent and meets all company standards. Moving to human approval."
	}

	def human_approval_node(state: WorkflowState):
	# This node doesn't actually do anything automatically.
	# It serves as a placeholder for the graph to pause.
	# When the human approves via API, we will just pass the state through.
	return {"status": "human_approved"}

	# Routing
	def route_after_critic(state: WorkflowState):
	if state["status"] == "approved":
	return "human_approval"
	elif state["revision_count"] >= 3:
	# Prevent infinite loops, force human review if too many revisions
	return "human_approval"
	else:
	return "writer"

	def build_workflow():
	workflow = StateGraph(WorkflowState)

	# Add nodes
	workflow.add_node("writer", writer_node)
	workflow.add_node("critic", critic_node)
	workflow.add_node("human_approval", human_approval_node)

	# Set entry point
	workflow.set_entry_point("writer")

	# Add edges
	workflow.add_edge("writer", "critic")
	workflow.add_conditional_edges(
	"critic",
	route_after_critic,
	{
	"writer": "writer",
	"human_approval": "human_approval"
	}
	)
	workflow.add_edge("human_approval", END)

	# Compile graph with a checkpointer for human-in-the-loop (memory)
	from langgraph.checkpoint.memory import MemorySaver
	memory = MemorySaver()

	# The graph will interrupt BEFORE the human_approval node
	app = workflow.compile(
	checkpointer=memory,
	interrupt_before=["human_approval"]
	)

	return app

	# Singleton instance of the graph
	workflow_app = build_workflow()