from typing import TypedDict, Annotated

from langgraph.graph import StateGraph, END

from langchain_core.messages import BaseMessage

import operator



# 상태 정의

class GraphState(TypedDict):

    messages: Annotated[list[BaseMessage], operator.add]

    # 추가 상태 필드

    user_input: str

    context: str



# 노드 정의

def retrieve_node(state):

    # 검색 로직

    return {"messages": [message]}



def generate_node(state):

    # 생성 로직

    return {"messages": [message]}



# 그래프 생성

workflow = StateGraph(GraphState)

workflow.add_node("retrieve", retrieve_node)

workflow.add_node("generate", generate_node)

workflow.add_edge("retrieve", "generate")

workflow.set_entry_point("retrieve")

workflow.set_finish_point("generate")

from langchain_core.prompts import PromptTemplate

from langchain_core.output_parsers import StrOutputParser

from langchain_openai import OpenAIEmbeddings, ChatOpenAI

from langchain_community.vectorstores import Chroma

from langchain_core.runnables import RunnableLambda

import operator



# RAG 상태 정의

class RagState(TypedDict):

    question: str

    context: str

    answer: str

    validation: bool



# 검색 노드

def retrieve_document(state):

    embeddings = OpenAIEmbeddings()

    vectorstore = Chroma(persist_directory="./chroma_db", embedding_function=embeddings)

    docs = vectorstore.similarity_search(state["question"], k=3)

    context = "\n".join([doc.page_content for doc in docs])

    return {"context": context}



# 생성 노드

def generate_answer(state):

    prompt = PromptTemplate.from_template(

        "Context: {context}\n\nQuestion: {question}\n\nAnswer:"

    )

    llm = ChatOpenAI(model="gpt-4", temperature=0.1)

    chain = prompt | llm | StrOutputParser()

    answer = chain.invoke({"context": state["context"], "question": state["question"]})

    return {"answer": answer}



# 검증 노드

def validate_answer(state):

    prompt = PromptTemplate.from_template(

        "Validate if the answer is correct for the question.\n\nQuestion: {question}\n\nAnswer: {answer}\n\nAnswer (True/False):"

    )

    llm = ChatOpenAI(model="gpt-4", temperature=0.1)

    chain = prompt | llm | StrOutputParser()

    validation = chain.invoke({"question": state["question"], "answer": state["answer"]})

    return {"validation": validation.lower() == "true"}



# 전체 워크플로우

def create_rag_workflow():

    workflow = StateGraph(RagState)



    workflow.add_node("retrieve", retrieve_document)

    workflow.add_node("generate", generate_answer)

    workflow.add_node("validate", validate_answer)



    workflow.add_edge("retrieve", "generate")

    workflow.add_edge("generate", "validate")



    workflow.set_entry_point("retrieve")

    workflow.set_finish_point("validate")



    return workflow.compile()



# 사용 예시

rag_workflow = create_rag_workflow()

result = rag_workflow.invoke({"question": "Python의 가상환경 생성 방법은?"})

from typing import List, Dict, Any

from langchain.tools import tool

from langchain_core.messages import ToolMessage



class MultiToolState(TypedDict):

    tasks: List[str]

    current_task: str

    execution_result: str

    feedback: str



# 도구 정의

@tool

def search_web(query: str) -> str:

    """웹에서 정보 검색"""

    # 실제 구현은 web search API 호출

    return f"Search results for '{query}'"



@tool

def calculate_math(expression: str) -> str:

    """수학 계산"""

    # 실제 구현은 계산기 로직

    return f"Result of {expression}"



@tool

def get_weather(city: str) -> str:

    """날씨 정보 가져오기"""

    # 실제 구현은 날씨 API 호출

    return f"Weather in {city}: Sunny, 25°C"



# 계획 노드

def plan_task(state):

    # 현재 작업 계획 생성

    tasks = [

        "Research technical specifications",

        "Perform mathematical calculations",

        "Check current weather conditions"

    ]

    return {"tasks": tasks, "current_task": tasks[0]}



# 실행 노드

def execute_task(state):

    # 현재 작업 실행

    tool_map = {

        "search": search_web,

        "calculate": calculate_math,

        "weather": get_weather

    }



    # 실제 실행 로직 (간소화됨)

    task = state["current_task"]

    result = f"Executed task: {task}"

    return {"execution_result": result}



# 관찰 노드

def observe_result(state):

    # 실행 결과 분석

    return {"feedback": f"Task '{state['current_task']}' completed successfully"}



# 결정 노드

def decide_next(state):

    # 다음 작업 결정

    current_index = state["tasks"].index(state["current_task"])

    if current_index < len(state["tasks"]) - 1:

        return {"current_task": state["tasks"][current_index + 1]}

    return {"current_task": None}



# 전체 워크플로우

def create_multitool_workflow():

    workflow = StateGraph(MultiToolState)



    workflow.add_node("plan", plan_task)

    workflow.add_node("execute", execute_task)

    workflow.add_node("observe", observe_result)

    workflow.add_node("decide", decide_next)



    workflow.add_edge("plan", "execute")

    workflow.add_edge("execute", "observe")

    workflow.add_edge("observe", "decide")

    workflow.add_edge("decide", "execute")



    workflow.set_entry_point("plan")

    workflow.set_finish_point("decide")



    return workflow.compile()

from typing import Optional

import asyncio



class HumanInLoopState(TypedDict):

    task_data: str

    human_review: Optional[str]

    approval: Optional[bool]

    processed_data: str



# 작업 생성 노드

def create_task(state):

    return {"task_data": "Data processing required"}



# 일시정지 노드

def pause_for_review(state):

    # 작업을 일시정지하고 인간 검토를 요청

    print("Waiting for human review...")

    return {"human_review": "Please review task data"}



# 검토 노드

async def review_task(state):

    # 인간 검토 기다리기 (실제 구현에서는 API 또는 사용자 인터페이스)

    await asyncio.sleep(1)  # 실제 구현에서는 사용자 입력 대기

    return {"human_review": "Review completed"}



# 진행 노드

def continue_processing(state):

    # 인간 승인 여부에 따라 진행

    if state.get("approval") is True:

        return {"processed_data": f"Processed: {state['task_data']}"}

    return {"processed_data": "Failed: Not approved"}



# 인간-함께-작업 워크플로우

def create_human_loop_workflow():

    workflow = StateGraph(HumanInLoopState)



    workflow.add_node("create", create_task)

    workflow.add_node("pause", pause_for_review)

    workflow.add_node("review", review_task)

    workflow.add_node("continue", continue_processing)



    workflow.add_edge("create", "pause")

    workflow.add_edge("pause", "review")

    workflow.add_edge("review", "continue")



    workflow.set_entry_point("create")

    workflow.set_finish_point("continue")



    return workflow.compile()