Refactor agent.py to use HuggingFaceEndpoint for model initialization and add AgentState TypedDict for message structure
Browse files
agent.py
CHANGED
|
@@ -4,8 +4,16 @@ from langgraph.graph import StateGraph, START, END
|
|
| 4 |
from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFacePipeline
|
| 5 |
from langchain_core.messages import HumanMessage
|
| 6 |
|
| 7 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 8 |
|
| 9 |
-
|
| 10 |
|
|
|
|
|
|
|
| 11 |
|
|
|
|
| 4 |
from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFacePipeline
|
| 5 |
from langchain_core.messages import HumanMessage
|
| 6 |
|
| 7 |
+
# Base Hugging Face LLM used by the chat wrapper
|
| 8 |
+
base_llm = HuggingFaceEndpoint(
|
| 9 |
+
repo_id="moonshotai/Kimi-K2.5",
|
| 10 |
+
task="text-generation",
|
| 11 |
+
temperature=0.0,
|
| 12 |
+
huggingfacehub_api_token=os.getenv("HUGGINGFACEHUB_API_TOKEN"),
|
| 13 |
+
)
|
| 14 |
|
| 15 |
+
model = ChatHuggingFace(llm=base_llm)
|
| 16 |
|
| 17 |
+
class AgentState(TypedDict):
|
| 18 |
+
messages: List[Dict[str, Any]]
|
| 19 |
|