chilicareAI / src /chains /rag.py
feryms's picture
big update
a13a62d
import os
from dotenv import load_dotenv
from langchain_groq import ChatGroq
from langchain_openai import ChatOpenAI
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.retrievers import BM25Retriever
from langchain_classic.retrievers import EnsembleRetriever
from langchain_core.documents import Document
from src.chains.prompt import DISEASE_PROMPT_TEMPLATE
load_dotenv()
llm = ChatGroq(
model="openai/gpt-oss-20b",
temperature=0.2,
api_key=os.getenv("GROQ_API_KEY"),
)
print("Memuat koneksi ke Database...")
embeddings = HuggingFaceEmbeddings(model_name="Qwen/Qwen3-Embedding-0.6B")
vectorstore = Chroma(
persist_directory="chroma_data",
embedding_function=embeddings,
collection_name="chilicare_kb"
)
chain = DISEASE_PROMPT_TEMPLATE | llm
def generate_narrative(disease_name):
print(f"Mencari data untuk label: {disease_name}...")
search_query = f"Penjelasan lengkap mengenai penyebab, ciri-ciri gejala, dan cara mengatasi penyakit {disease_name} pada tanaman cabai."
results = vectorstore.similarity_search(
query=search_query,
k=3, # Mengambil 3 potongan (chunks) teratas
filter={"label": disease_name}
)
if not results:
return f"Data penyakit '{disease_name}' tidak ditemukan di database."
retrieved_context = "\n\n".join([doc.page_content for doc in results])
print("Data ditemukan. Menghasilkan narasi dengan LLM...")
response = chain.invoke({
"disease_name": disease_name,
"context": retrieved_context
})
return response.content