cookidooAI/connect_chaint.py

70 lines
2.1 KiB
Python
Executable File

# 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")