# Plik: test_connect_chaint.py
import sys
import warnings
from langchain_openai import ChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate

warnings.filterwarnings("ignore")

# 1. Załaduj bazę (Tę, którą zbudował EpubToRag)
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vectorstore = Chroma(persist_directory="/home/pali112/db_build123d_reference", embedding_function=embeddings)

# 2. Połącz się z serwerem Llama (Qwen)
llm = ChatOpenAI(
    base_url="http://localhost:8081/v1",
    api_key="brak",
    model="local-model",
    temperature=0
)

# 3. System Prompt
template = """Jesteś ekspertem od biblioteki build123d.
Użyj poniższych fragmentów dokumentacji, aby odpowiedzieć na pytanie.
Jeśli nie znasz odpowiedzi, napisz, że nie wiesz.


WZORZEC KODU (Naśladuj ten styl):
--------------------------------------------------
from build123d import *

# 1. Definicja parametrów
length, width = 100, 50

# 2. Geometria w bloku 'with' (Builder Mode)
with BuildPart() as my_part:
    Box(length, width, 10)
    # Wycinanie otworu (Mode.SUBTRACT)
    with Locations(Location((0, 0, 0))):
        Cylinder(radius=5, height=10, mode=Mode.SUBTRACT)

# 3. Poprawny eksport
export_stl(my_part.part, "wynik.stl")

Dokumentacja:
{context}

Pytanie: {question}
Odpowiedź eksperta:"""

PROMPT = PromptTemplate(template=template, input_variables=["context", "question"])

# 4. Konfiguracja RAG (To jest 'qa_chain', którego szuka most_ai.py)
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=vectorstore.as_retriever(search_kwargs={"k": 5}),
    return_source_documents=True,
    chain_type_kwargs={"prompt": PROMPT}
)

# Zabezpieczenie przed przypadkowym uruchomieniem
if __name__ == "__main__":
    if len(sys.argv) > 1:
        print(qa_chain.invoke(" ".join(sys.argv[1:]))["result"])
    else:
        print("To jest moduł backendu. Uruchom most_ai.py")