reddwarf03/archethic-sdk-assistant

1import logging2import sys3import chromadb4import gradio as gr5import os6 7from llama_index.core import VectorStoreIndex, Settings, get_response_synthesizer8from llama_index.core.embeddings import resolve_embed_model9from llama_index.vector_stores.chroma import ChromaVectorStore10from llama_index.llms.huggingface import HuggingFaceLLM11from llama_index.core.response_synthesizers import ResponseMode12 13from pathlib import Path14 15# Setup logging16logging.basicConfig(stream=sys.stdout, level=logging.INFO)17log = logging.getLogger(__name__)18log.addHandler(logging.StreamHandler(stream=sys.stdout))19 20# --- Configuration --- 21CHROMA_DB_PATH = Path("./vector_store") # Path relative to app.py inside hf_space22COLLECTION_NAME = "archethic_sdk_docs"23EMBED_MODEL_NAME = "local:BAAI/bge-small-en-v1.5"24# LLM Configuration25LLM_MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.2"26LLM_MAX_NEW_TOKENS = 1024 # Max tokens for the LLM to generate27LLM_CONTEXT_WINDOW = 3900 # Context window should be okay for LLM version too28 29# System Prompt for the LLM30SYSTEM_PROMPT = """You are an expert Archethic blockchain assistant.31Your goal is to answer questions about developing on Archethic using the provided context.32If the question asks for code (JavaScript or AssemblyScript), generate functional code based ONLY on the context.33If the context does not contain the answer or enough information to generate the code, state that clearly.34Be concise and helpful.35"""36 37# --- Global Variables --- 38query_engine = None39 40def initialize_rag_engine():41    """Initializes the RAG query engine using ChromaDB and HuggingFaceLLM."""42    global query_engine43    if query_engine is not None:44        log.info("RAG engine already initialized.")45        return46 47    log.info("--- Initializing RAG Engine ---")48 49    # --- Embedding Model ---50    log.info(f"Setting up embedding model: {EMBED_MODEL_NAME}")51    embed_model = resolve_embed_model(EMBED_MODEL_NAME)52    Settings.embed_model = embed_model53 54    # --- Hugging Face Auth ---55    hf_token = os.getenv("HUGGINGFACE_TOKEN")56    if not hf_token:57        log.warning("HUGGINGFACE_TOKEN secret not found. Inference API calls may fail or be rate-limited.")58    else:59        log.info("HUGGINGFACE_TOKEN found and will be used automatically by huggingface_hub.")60        from huggingface_hub import login61        login(hf_token)62 63    # --- LLM Configuration (HuggingFace Inference API) ---64    log.info(f"Setting up HuggingFaceLLM with model: {LLM_MODEL_ID}")65    66    llm = HuggingFaceLLM(67        model_name=LLM_MODEL_ID,68        tokenizer_name=LLM_MODEL_ID,69        context_window=LLM_CONTEXT_WINDOW,70        max_new_tokens=LLM_MAX_NEW_TOKENS,71        generate_kwargs={"temperature": 0.7, "do_sample": True},72    )73    Settings.llm = llm74    75    # --- Load Vector Store ---76    log.info(f"Loading existing ChromaDB client from: {CHROMA_DB_PATH}")77    try:78        db = chromadb.PersistentClient(path=str(CHROMA_DB_PATH))79        chroma_collection = db.get_collection(COLLECTION_NAME)80        vector_store = ChromaVectorStore(chroma_collection=chroma_collection)81        index = VectorStoreIndex.from_vector_store(82            vector_store, 83            embed_model=embed_model84        )85        log.info("Vector store loaded successfully.")86    except Exception as e:87        log.error(f"Error loading vector store: {e}", exc_info=True)88        raise89 90    # --- Response Synthesizer ---91    log.info("Creating Response Synthesizer with streaming enabled...")92    response_synthesizer = get_response_synthesizer(93        response_mode=ResponseMode.COMPACT,94        llm=llm,95        streaming=True,96    )97    98    # --- Query Engine ---99    log.info("Creating query engine with LLM and custom synthesizer...")100    query_engine = index.as_query_engine(101        similarity_top_k=3, 102        response_synthesizer=response_synthesizer,103    )104    log.info(f"Query engine created. Type: {type(query_engine)}") 105 106    log.info("RAG engine initialized successfully.")107    log.info("---------------------------------")108 109 110 111# Modified function to stream the response112def answer_query(query_text: str):113    """Answers a query using the initialized RAG engine and streams the response."""114    if not query_engine:115        log.error("Query engine is not initialized. Call initialize_rag_engine() first.")116        yield "Error: RAG Engine not ready."117        return118 119    log.info(f"Received query for streaming: {query_text}")120        121    try:122        # Use the query method with the user's text123        streaming_response = query_engine.query(query_text)124 125        # Check for streaming attribute126        if not hasattr(streaming_response, 'response_gen'):127            log.error("Response object does not support streaming (no 'response_gen').")128            yield "Error: Streaming not supported by the current response object."129            if hasattr(streaming_response, 'response'): # Fallback130                yield str(streaming_response.response)131            return132 133        # Stream the response tokens134        partial_response = ""135        for token in streaming_response.response_gen:136            partial_response += token137            yield partial_response 138        log.info("Streaming finished.")139 140    except Exception as e:141        log.error(f"Error during query execution: {e}", exc_info=True)142        yield f"Error processing query: {e}"143 144 145# --- Gradio Interface Setup --- 146if __name__ == "__main__":147    log.info("Starting Gradio app initialization...")148    149    try:150        initialize_rag_engine()151    except Exception as e:152        log.error(f"Fatal error initializing RAG engine: {e}", exc_info=True)153        sys.exit("Failed to initialize RAG engine. Check logs.")154 155    log.info("Setting up Gradio interface...")156    iface = gr.Interface(157        fn=answer_query,158        inputs=gr.Textbox(lines=7, label="Your Question:", placeholder="Ask anything about Archethic SDKs or Docs..."),159        outputs=gr.Textbox(label="AI Assistant Answer", interactive=False, show_copy_button=True),160        title="🤖 Archethic SDK AI Assistant",161        description="Ask questions about Archethic development (JS SDK, AssemblyScript Contracts, Docs). The AI will use the documentation and code examples as context. Uses Mistral-7B-Instruct via HF Inference API.",162        allow_flagging="never",163    )164 165    log.info("Launching Gradio interface...")166    iface.launch(server_name="0.0.0.0") 167    log.info("Gradio app launched.")