from langchain_deepseek import ChatDeepSeek
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel, Field
from typing import Literal


class RouteQuery(BaseModel):
    """
    Route a user query to the most relevant datasource
    """

    datasource: Literal["vectorstore","websearch"] = Field(
        ...,
        description="Given a user question choose to route it to websearch or a vectorstore"
    )

llm = ChatDeepSeek(model="deepseek-chat", temperature=0)
structured_llm_router = llm.with_structured_output(RouteQuery)

system_prompt = """ You are an expert at routing a user question to a vectorstore or web search.\n
The vectorstore contains documents related to agents, prompt engineering and adversarial attacks.\
Use the vectorstore for questions on these topics. For all else, use websearch."""

route_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system_prompt),
        ("human", "{question}")
    ]
)

question_router = route_prompt | structured_llm_router