Untitled

from langchain_core.output_parsers import JsonOutputParser
from pydantic import BaseModel, Field
from langchain_core.messages import HumanMessage
from langchain_google_vertexai import ChatVertexAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from typing import List

llm = ChatVertexAI(
    model_name="gemini-1.5-flash-001", 
    temperature=0.5, 
    max_output_tokens=1024,
    response_mime_type="application/json",
    
)


template = ChatPromptTemplate.from_messages(
    [
        ("system", """You are a search engine trying to understand what questions will be asked from each of the knowledge base articles you have access to."""),
        MessagesPlaceholder(variable_name="messages")
    ]
)



class QueryList(BaseModel):
    queries: List[str] = Field(description="A list of queries that can be answered by the knowledge base article'")

parser = JsonOutputParser(pydantic_object=QueryList)

template.partial_variables = {"format_instructions": parser.get_format_instructions()}
chain = template | llm | parser

def send_llm_request(knowledge_base_text):
    json_example = {
      "queries": [
        "How do I reset my password?",
        "What's the number of the concat center?"
      ]
    }

    prompt = f"""What are 1-5 questions a user might have that would be answered by this article and only this article?
    
    =================
    Article:    
    {knowledge_base_text}
    
    ================
    Return a json format like this and double check that the questions can be answered just with the text provided in the article:
    ```
    {str(json_example)}
    ```
    """
    parts = []
    parts.append(
        {
            "type": "text",
            "text": prompt
        }
    )
    messages = [HumanMessage(content=parts).dict()]
    json_text = chain.invoke(messages)
    return json_text['queries']

# Test the prompt
send_llm_request("This article talks about Vertex AI")