#pip install -U langchain-ollama

from sentence_transformers import SentenceTransformer
#from langchain.vectorstores import FAISS
from langchain_community.vectorstores import FAISS
from langchain.embeddings.base import Embeddings
from langchain.schema import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.chat_models import ChatOllama
import requests

# --- 1. Custom Embedding Class για LangChain wrapper ---
print("---STEP 1: Embedding Class για LangChain wrapper---")
class SentenceTransformerEmbeddings(Embeddings):
    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        self.model = SentenceTransformer(model_name)

    def embed_documents(self, texts):
        return self.model.encode(texts, convert_to_tensor=False).tolist()

    def embed_query(self, text):
        return self.model.encode(text, convert_to_tensor=False).tolist()

# --- 2. Προετοιμασία Εγγράφων ---
print("---STEP 2: read texts from large_output.md---")
with open('large_output.md', 'r', encoding='utf-8') as file:
    raw_texts = [line.strip() for line in file.readlines()]
docs = [Document(page_content=text) for text in raw_texts]

# --- 3. Embeddings & FAISS VectorStore ---
print("---STEP 3: Embeddings & FAISS VectorStore---")
embeddings = SentenceTransformerEmbeddings()
vectorstore = FAISS.from_documents(docs, embedding=embeddings)

# --- 4. LLM & RAG Chain ---
print("---STEP 4: read queries from queries_all.txt and ask LLM---")
count = 0
with open('queries_all.txt', 'r', encoding='utf-8') as f:
    for line in f:
        query = line.strip()
        results_with_scores = vectorstore.similarity_search_with_score(query, k=3)
        context = "\n---\n".join(
            f"{doc.page_content}"
            for doc, _ in results_with_scores
        )

        prompt = ChatPromptTemplate.from_messages([
            ("system",
             "Answer the following question based on the provided context."
             "Context:\n{context}\n"),
            ("user", "{question}")
        ])

        llm = ChatOllama(
            model="deepseek-r1:70b",
            temperature=0.1,
            base_url="http://195.130.94.63:11434"
        )

        chain = prompt | llm
        response = chain.invoke({"question": query, "context": context})

        with open('responses_all_ollama_model_deepseek-r1-70b_w_rag.txt', 'a', encoding='utf-8') as file:
            count += 1
            file.write("Prompt " + str(count) + " :\n" +
                "---Question: " + str(query) +
                "\n---Context: " + str(context) +
                "\n\n---Anwer:\n" +
                str(response.content) + "\n\n\n")
            print(count)