Update app.py

app.py

@@ -1,89 +1,119 @@
-import os
 import streamlit as st
-from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
+import os
+from llama_index.core import (
+    VectorStoreIndex,
+    SimpleDirectoryReader,
+    Settings,
+)
+from llama_index.core import PromptTemplate
 from llama_index.llms.gemini import Gemini
-from llama_index.core.extractors import TitleExtractor
-from llama_index.core.node_parser import SentenceWindowNodeParser
-from llama_index.core.ingestion import IngestionPipeline
-from llama_index.core.query_engine import RetrieverQueryEngine
-from llama_index.core.retrievers import AutoMergingRetriever
-from llama_index.core.indices.vector_store.retrievers import VectorIndexRetriever
-from llama_index.vector_stores.chroma import ChromaVectorStore
-
-import chromadb
-
+from llama_index.embeddings.gemini import GeminiEmbedding
+import logging
+import google.generativeai as genai
 from dotenv import load_dotenv
+from pathlib import Path
 
 load_dotenv()
 
-# Configurer les paramètres globaux
-Settings.llm = Gemini(api_key=os.environ["GOOGLE_API_KEY"], temperature=0.1)
-Settings.chunk_size = 1024  # Taille des chunks pour l'indexation
-# Nombre de tokens générés par le LLM
-
-# Fonction pour charger les données et créer l'index (optimisé pour éviter les rechargements inutiles)
-@st.cache_resource
-def load_data_and_create_index():
-    """Charge les documents PDF et crée l'index vectoriel."""
-    documents = SimpleDirectoryReader("./data").load_data()
-
-    # Créer un pipeline d'ingestion avec extraction de titre et fenêtrage de phrases
-    node_parser = SentenceWindowNodeParser.from_defaults(
-        window_size=3,
-        window_metadata_key="window",
-        original_text_metadata_key="original_text",
-    )
-
-    text_splitter = node_parser.get_leaf_nodes_and_parent_nodes
-    extractors = [TitleExtractor(nodes=5)]
+# Set logging level
+logging.basicConfig(level=logging.INFO)
 
-    pipeline = IngestionPipeline(
-        transformations=[node_parser, *extractors]
-    )
-
-    # Indexer les documents
-    nodes = pipeline.run(documents=documents)
-
-    # Initialiser la base de données vectorielle (exemple avec Chroma)
-    db = chromadb.Client()
-    chroma_collection = db.get_or_create_collection("legal_docs")
-    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
-
-    # Créer l'index
-    index = VectorStoreIndex.from_documents(nodes, vector_store=vector_store)
-    return index
-
-# Fonction pour effectuer la requête
-def perform_query(query_str, index):
-    """Effectue une requête sur l'index et renvoie la réponse."""
-    # Créer un AutoMergingRetriever
-    base_retriever = VectorIndexRetriever(
-        index=index,
-        similarity_top_k=8,
-    )
-    retriever = AutoMergingRetriever(base_retriever, index.storage_context)
 
-    # Créer le moteur de requête
-    query_engine = RetrieverQueryEngine.from_args(retriever=retriever)
+# Configure Gemini Pro
+genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
 
-    response = query_engine.query(query_str)
-    return response
+model_gemini_pro_vision = "gemini-pro-vision"
 
-# Interface utilisateur Streamlit
-st.title("Agent de Questions-Réponses Juridiques (Gemini + LlamaIndex)")
+# Configure Gemini models
+Settings.llm = Gemini(model=model_gemini_pro_vision,
+                        api_key=os.getenv("GOOGLE_API_KEY"))
+Settings.embed_model = GeminiEmbedding(
+    model_name="models/embedding-001",
+                        api_key=os.getenv("GOOGLE_API_KEY")
+    )
 
-# Charger les données et créer l'index (une seule fois)
-index = load_data_and_create_index()
+def load_and_index_pdf(pdf_path):
+    """Loads and index the pdf.
+    
+    Args : 
+    pdf_path (str) : The path to the pdf file
+    
+    Returns : 
+    index (llama_index.core.VectorStoreIndex): The vector index
+    """
+    try:
+        logging.info(f"Loading PDF document from: {pdf_path}")
+        documents = SimpleDirectoryReader(input_files=[pdf_path]).load_data()
+        if documents:
+            logging.info("Creating vector store index")
+            index = VectorStoreIndex.from_documents(documents)
+            return index
+        else:
+            logging.warning("No documents found in the PDF")
+            return None
+    except Exception as e:
+        logging.error(f"Error loading and indexing PDF: {e}")
+        return None
+    
+def translate_text(french_text, index):
+    """Translates french text to Yipunu.
+    
+    Args : 
+    french_text (str): The french text to translate.
+    index (llama_index.core.VectorStoreIndex): The vector index.
+    
+    Returns:
+     (str): The yipunu translation or an error message.
+    """
+
+    try:
+        logging.info(f"Initiating translation of: {french_text}")
+
+        template = (
+            "Tu es un excellent traducteur du français vers le yipunu. Tu traduis le texte sans donner d'explication. "
+            "Texte: {french_text} "
+            "Traduction:"
+        )
+
+        prompt_template = PromptTemplate(template)
+
+        query_engine = index.as_query_engine(
+            text_qa_template=prompt_template
+        )
+        response = query_engine.query(french_text)
+        logging.info(f"Translation Result: {response.response}")
+        return response.response
+    except Exception as e:
+        logging.error(f"Error during translation: {e}")
+        return f"Error during translation: {str(e)}"
+        
+
+def main():
+    """Main function for streamlit app."""
+    
+    st.title("French to Yipunu Translation App")
+
+    # PDF File Upload
+    uploaded_file = st.file_uploader("Upload a PDF file containing the Punu grammar:", type="pdf")
+
+    if uploaded_file is not None:
+            # Save file to a temporary location
+        temp_file_path = Path("temp_file.pdf")
+        with open(temp_file_path, "wb") as f:
+            f.write(uploaded_file.read())
+
+        index = load_and_index_pdf(str(temp_file_path))
+        if index:
+            french_text = st.text_area("Enter French Text:", "Ni vosi yipunu")
+            if st.button("Translate"):
+                 translation = translate_text(french_text, index)
+                 st.success(f"Yipunu Translation: {translation}")
+
+        # Clean up temp files
+        os.remove(temp_file_path)
+    else:
+        st.info("Please upload a pdf containing the punu grammar.")
 
-# Champ de saisie de la question
-query_str = st.text_input("Posez votre question juridique ici :")
 
-# Bouton pour soumettre la question
-if st.button("Poser la question"):
-    if query_str:
-        with st.spinner("Recherche en cours..."):
-            response = perform_query(query_str, index)
-        st.success("Réponse :")
-        st.write(response)
-    else:
-        st.error("Veuillez saisir une question.")
+if __name__ == "__main__":
+    main()