|
|
|
|
|
|
|
|
from langgraph.graph import StateGraph, START, END |
|
|
from typing import TypedDict, Annotated, List, Generator, Union, Dict, Any |
|
|
from dotenv import load_dotenv |
|
|
from langgraph.checkpoint.memory import InMemorySaver |
|
|
from langchain_google_genai import ChatGoogleGenerativeAI |
|
|
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, AIMessageChunk |
|
|
from langgraph.graph.message import add_messages |
|
|
import logging |
|
|
|
|
|
|
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
load_dotenv() |
|
|
THREAD_ID = "1" |
|
|
|
|
|
|
|
|
|
|
|
gen_llm = ChatGoogleGenerativeAI(model='gemini-1.5-flash') |
|
|
|
|
|
class ChatState(TypedDict): |
|
|
messages: Annotated[list[BaseMessage], add_messages] |
|
|
|
|
|
def chat_node(state: ChatState): |
|
|
messages = state['messages'] |
|
|
response = gen_llm.invoke(messages) |
|
|
return {"messages": [response]} |
|
|
|
|
|
|
|
|
checkpointer = InMemorySaver() |
|
|
|
|
|
graph = StateGraph(ChatState) |
|
|
graph.add_node("chat_node", chat_node) |
|
|
graph.add_edge(START, "chat_node") |
|
|
graph.add_edge("chat_node", END) |
|
|
|
|
|
chatbot = graph.compile(checkpointer=checkpointer) |
|
|
|
|
|
|
|
|
|