File size: 1,654 Bytes
6064cea 5f6d20c 6064cea 5f6d20c 6064cea 5f6d20c 6064cea 5f6d20c 6064cea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | from retrival.usage import LegalRetriever
from LLM.llm import LegalLLM
import os
import dotenv
dotenv.load_dotenv()
class RAGPipeline:
def __init__(self):
self.retriever = LegalRetriever()
self.llm = LegalLLM(api_key=os.getenv("GEMINI_API_KEY"))
# =====================================================
# CONTEXT
# =====================================================
def get_context(
self,
query: str
):
retrieval_result = (
self.retriever.retrieve(
query
)
)
return retrieval_result
# =====================================================
# LLM
# =====================================================
def get_llm_result(
self,
query: str,
context: str
):
return self.llm.generate(
query=query,
context=context
)
# =====================================================
# SEARCH
# =====================================================
def search(
self,
query: str
):
retrieval_result = (
self.get_context(
query
)
)
context = retrieval_result[
"final_context"
]
print("=" * 100)
print(retrieval_result["final_context"])
print("=" * 100)
answer = (
self.get_llm_result(
query,
context
)
)
return {
"query": query,
"answer": answer,
"retrieval": retrieval_result
} |