main.py

# app.py
from flask import Flask, request, jsonify
import asyncio
import os
from schemas import Task
from langchain.prompts import ChatPromptTemplate
from langchain.schema import HumanMessage, SystemMessage
from config import LANGCHAIN_TRACING_V2, chat_model
from task_manager import TaskManager
from orchestration import a_generate_task_tree
from tree_utils import print_task_tree
from evaluation import evaluate_task_decomposition
from langchain_core.tracers.context import tracing_v2_enabled
import json
from typing import Dict
from jsonschema import validate, ValidationError

from flask import Flask, request, jsonify
from flask_cors import CORS

app = Flask(__name__)


CORS(app)  # Enable CORS for all routes and origins

async def generate_task_tree(prompt: str):
    task_manager = TaskManager() #Creating task manager object here
    with tracing_v2_enabled(project_name="Task Decomposition") if LANGCHAIN_TRACING_V2 else open(os.devnull, "w") as f: # only trace if the relevant flag is turned on
        # TODO: How can I transform the user prompt to be more specific and actionable for the LLM?
        # system_message = SystemMessage(
        #     content="You are a computer use agent capable of doing anything. Rephrase the user's task prompt to highlight the key action verbs in the user's request and identify what needs to be done.")
        # human_message = HumanMessage(
        #     content=f"Output a very concise task prompt to help an LLM understand the user's task prompt: {prompt}.\n\nEmphasize what action verbs are specified by the user. Only output the prompt and nothing else.")

        # chat_prompt = ChatPromptTemplate.from_messages([system_message, human_message])

        # response = await chat_model.ainvoke(chat_prompt.format_messages())
        # response_content = response.content
        # print(f"{response_content=}")
        response_content = prompt

        full_task, _ = await a_generate_task_tree(response_content, Task.model_json_schema(), task_manager) # Passing task_manager object

        if full_task and validate_task(full_task, Task.model_json_schema()):
            # Save full_task to out.txt
            with open("out.txt", 'w') as file:
                json.dump(full_task, file, indent=4)

            return full_task
        else:
            return {"error": "Task generation or validation failed."}

def validate_task(task: Dict, schema: Dict): #Keeping this function here since it is tiny
    try:
        validate(instance=task, schema=schema)
        return True
    except ValidationError as e:
        print(f"Task validation error: {e}")
        return False


# Your routes here
@app.route('/generate_task_tree', methods=['POST'])
def generate_task_tree_api():
    data = request.get_json()
    if 'prompt' not in data:
        return jsonify({"error": "Missing 'prompt' in request"}), 400
    
    prompt = data['prompt']
    prompt = prompt + ".Decompose the main task into atleast 2 subtasks and 1 ingest and produces."
    print(prompt)
    # Run the asynchronous function in a synchronous context
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    result = loop.run_until_complete(generate_task_tree(prompt))
    loop.close()
    
    return jsonify(result)

if __name__ == "__main__":
    app.run(host="127.0.0.1", port=5001)