from langchain.output_parsers import StructuredOutputParser, ResponseSchema
from llama_index.core.output_parsers import LangchainOutputParser
from llama_index.llms.openai import OpenAI
from llama_index.core.schema import TextNode
from llama_index.core import VectorStoreIndex
from pydantic import BaseModel
from typing import List

nodes = [
    TextNode(
        text="Róże mają żywe kolory i ładnie pachną."),
    TextNode(
        text="Dęby to wysokie drzewa zielonymi liśćmi."),
]

schemas = [
    ResponseSchema(
        name="odpowiedź", 
        description=(
            "odpowiedź na pytanie użytkownika"
        )
    ),
    ResponseSchema(
        name="źródło",
        description=(
            "tekst źródłowy użyty do opowiedzi użytkownikowi, "
            "powinien być cytatem z oryginalnego promptu."
        )
    )
]

lc_parser = StructuredOutputParser.from_response_schemas(schemas)
output_parser = LangchainOutputParser(lc_parser)

llm = OpenAI(output_parser=output_parser)

index = VectorStoreIndex(nodes=nodes)
query_engine = index.as_query_engine(llm=llm)
response = query_engine.query(
    "Czy dęby są niskie? tak czy nie",
)
print(response)
