Add RAG-based chatbot for querying ingested data

requirements.txt

@@ -6,3 +6,7 @@ tiktoken
 pytest
 pytest-asyncio
 click>=8.0.0
+sentence-transformers
+openai
+langchain
+faiss-cpu

src/gitingest/embedding.py

@@ -0,0 +1,25 @@
+from sentence_transformers import SentenceTransformer
+import numpy as np
+import os
+
+class EmbeddingHandler:
+    def __init__(self, model_name='all-MiniLM-L6-v2'):
+        self.model = SentenceTransformer(model_name)
+
+    def get_embeddings(self, texts):
+        return self.model.encode(texts, convert_to_tensor=True)
+
+    def save_embeddings(self, embeddings, file_path):
+        np.save(file_path, embeddings)
+
+    def load_embeddings(self, file_path):
+        return np.load(file_path)
+
+    def vectorize_repository(self, repo_path):
+        file_contents = []
+        for root, _, files in os.walk(repo_path):
+            for file in files:
+                file_path = os.path.join(root, file)
+                with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
+                    file_contents.append(f.read())
+        return self.get_embeddings(file_contents)

src/gitingest/rag_chatbot.py

@@ -0,0 +1,41 @@
+import os
+import openai
+from langchain.chains import RetrievalQA
+from langchain.llms import OpenAI
+from langchain.vectorstores import FAISS
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.prompts import PromptTemplate
+
+class RAGChatbot:
+    def __init__(self, openai_api_key, vector_db_path):
+        self.openai_api_key = openai_api_key
+        openai.api_key = self.openai_api_key
+        self.vector_db_path = vector_db_path
+        self.vector_store = self.load_vector_store()
+
+    def load_vector_store(self):
+        if os.path.exists(self.vector_db_path):
+            return FAISS.load_local(self.vector_db_path, OpenAIEmbeddings())
+        else:
+            raise FileNotFoundError(f"Vector database not found at {self.vector_db_path}")
+
+    def save_vector_store(self):
+        self.vector_store.save_local(self.vector_db_path)
+
+    def update_vector_store(self, texts):
+        embeddings = OpenAIEmbeddings()
+        new_vectors = embeddings.embed_documents(texts)
+        self.vector_store.add_documents(new_vectors)
+        self.save_vector_store()
+
+    def generate_response(self, query):
+        retriever = self.vector_store.as_retriever()
+        qa_chain = RetrievalQA(
+            retriever=retriever,
+            llm=OpenAI(api_key=self.openai_api_key),
+            prompt_template=PromptTemplate(
+                input_variables=["context", "question"],
+                template="Context: {context}\n\nQuestion: {question}\n\nAnswer:"
+            )
+        )
+        return qa_chain.run(query)

src/gitingest/tests/test_chatbot_endpoints.py

@@ -0,0 +1,25 @@
+import pytest
+from fastapi.testclient import TestClient
+from main import app
+
+client = TestClient(app)
+
+@pytest.fixture
+def mock_generate_response(monkeypatch):
+    def mock_response(self, query):
+        return "Mock response"
+    monkeypatch.setattr("gitingest.rag_chatbot.RAGChatbot.generate_response", mock_response)
+
+def test_query_chatbot(mock_generate_response):
+    response = client.post("/chatbot/query", json={"query": "Test query"})
+    assert response.status_code == 200
+    assert response.json() == {"response": "Mock response"}
+
+def test_query_chatbot_error(mock_generate_response, monkeypatch):
+    def mock_response_error(self, query):
+        raise Exception("Mock error")
+    monkeypatch.setattr("gitingest.rag_chatbot.RAGChatbot.generate_response", mock_response_error)
+
+    response = client.post("/chatbot/query", json={"query": "Test query"})
+    assert response.status_code == 500
+    assert response.json() == {"detail": "Mock error"}

src/gitingest/tests/test_embedding.py

@@ -0,0 +1,33 @@
+import os
+import numpy as np
+import pytest
+from gitingest.embedding import EmbeddingHandler
+
+@pytest.fixture
+def embedding_handler():
+    return EmbeddingHandler()
+
+def test_get_embeddings(embedding_handler):
+    texts = ["Hello world", "Test sentence"]
+    embeddings = embedding_handler.get_embeddings(texts)
+    assert embeddings.shape[0] == 2
+    assert embeddings.shape[1] > 0
+
+def test_save_and_load_embeddings(embedding_handler, tmp_path):
+    texts = ["Hello world", "Test sentence"]
+    embeddings = embedding_handler.get_embeddings(texts)
+    file_path = os.path.join(tmp_path, "embeddings.npy")
+    embedding_handler.save_embeddings(embeddings, file_path)
+    loaded_embeddings = embedding_handler.load_embeddings(file_path)
+    assert np.array_equal(embeddings, loaded_embeddings)
+
+def test_vectorize_repository(embedding_handler, tmp_path):
+    repo_path = tmp_path / "repo"
+    repo_path.mkdir()
+    file1 = repo_path / "file1.txt"
+    file1.write_text("Hello world")
+    file2 = repo_path / "file2.txt"
+    file2.write_text("Test sentence")
+    embeddings = embedding_handler.vectorize_repository(str(repo_path))
+    assert embeddings.shape[0] == 2
+    assert embeddings.shape[1] > 0

src/gitingest/tests/test_rag_chatbot.py

@@ -0,0 +1,28 @@
+import pytest
+from gitingest.rag_chatbot import RAGChatbot
+
+@pytest.fixture
+def chatbot():
+    return RAGChatbot(openai_api_key="your_openai_api_key", vector_db_path="path_to_vector_db")
+
+def test_generate_response(chatbot):
+    query = "What is the purpose of this repository?"
+    response = chatbot.generate_response(query)
+    assert response is not None
+    assert isinstance(response, str)
+    assert len(response) > 0
+
+def test_update_vector_store(chatbot):
+    texts = ["This is a test document.", "Another test document."]
+    chatbot.update_vector_store(texts)
+    assert chatbot.vector_store is not None
+    assert len(chatbot.vector_store) > 0
+
+def test_load_vector_store(chatbot):
+    vector_store = chatbot.load_vector_store()
+    assert vector_store is not None
+    assert len(vector_store) > 0
+
+def test_save_vector_store(chatbot):
+    chatbot.save_vector_store()
+    assert chatbot.vector_db_path.exists()

src/routers/chatbot.py

@@ -0,0 +1,22 @@
+from fastapi import APIRouter, HTTPException
+from pydantic import BaseModel
+from gitingest.rag_chatbot import RAGChatbot
+
+router = APIRouter()
+
+class ChatbotRequest(BaseModel):
+    query: str
+
+class ChatbotResponse(BaseModel):
+    response: str
+
+# Initialize the RAGChatbot with your OpenAI API key and vector database path
+chatbot = RAGChatbot(openai_api_key="your_openai_api_key", vector_db_path="path_to_vector_db")
+
+@router.post("/chatbot/query", response_model=ChatbotResponse)
+async def query_chatbot(request: ChatbotRequest):
+    try:
+        response = chatbot.generate_response(request.query)
+        return ChatbotResponse(response=response)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))

src/templates/github.jinja

@@ -60,4 +60,100 @@
         });
     }
 </script>
-{% endblock extra_scripts %}
+
+<!-- Chatbot UI -->
+<div class="chatbot-container">
+    <div class="chatbot-header">
+        <h2>Chat with our RAG-based Bot</h2>
+    </div>
+    <div class="chatbot-messages" id="chatbot-messages"></div>
+    <div class="chatbot-input">
+        <input type="text" id="chatbot-input" placeholder="Type your message here...">
+        <button onclick="sendMessage()">Send</button>
+    </div>
+</div>
+
+<script>
+    async function sendMessage() {
+        const input = document.getElementById('chatbot-input');
+        const messages = document.getElementById('chatbot-messages');
+        const userMessage = input.value;
+        if (!userMessage) return;
+
+        // Display user message
+        const userMessageElement = document.createElement('div');
+        userMessageElement.className = 'user-message';
+        userMessageElement.textContent = userMessage;
+        messages.appendChild(userMessageElement);
+
+        // Clear input
+        input.value = '';
+
+        // Send message to chatbot API
+        try {
+            const response = await fetch('/chatbot/query', {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/json'
+                },
+                body: JSON.stringify({ query: userMessage })
+            });
+            const data = await response.json();
+
+            // Display chatbot response
+            const botMessageElement = document.createElement('div');
+            botMessageElement.className = 'bot-message';
+            botMessageElement.textContent = data.response;
+            messages.appendChild(botMessageElement);
+        } catch (error) {
+            console.error('Error:', error);
+        }
+    }
+</script>
+
+<style>
+    .chatbot-container {
+        border: 1px solid #ccc;
+        border-radius: 5px;
+        padding: 10px;
+        max-width: 400px;
+        margin: 20px auto;
+    }
+
+    .chatbot-header {
+        text-align: center;
+        font-weight: bold;
+    }
+
+    .chatbot-messages {
+        height: 200px;
+        overflow-y: auto;
+        border: 1px solid #ccc;
+        padding: 10px;
+        margin-bottom: 10px;
+    }
+
+    .chatbot-input {
+        display: flex;
+    }
+
+    .chatbot-input input {
+        flex: 1;
+        padding: 5px;
+    }
+
+    .chatbot-input button {
+        padding: 5px 10px;
+    }
+
+    .user-message {
+        text-align: right;
+        color: blue;
+    }
+
+    .bot-message {
+        text-align: left;
+        color: green;
+    }
+</style>
+{% endblock extra_scripts %}

src/templates/index.jinja

@@ -61,7 +61,100 @@
 
 {% include 'components/result.jinja' %}
 
+<!-- Chatbot UI -->
+<div class="chatbot-container">
+    <div class="chatbot-header">
+        <h2>Chat with our RAG-based Bot</h2>
+    </div>
+    <div class="chatbot-messages" id="chatbot-messages"></div>
+    <div class="chatbot-input">
+        <input type="text" id="chatbot-input" placeholder="Type your message here...">
+        <button onclick="sendMessage()">Send</button>
+    </div>
+</div>
+
+<script>
+    async function sendMessage() {
+        const input = document.getElementById('chatbot-input');
+        const messages = document.getElementById('chatbot-messages');
+        const userMessage = input.value;
+        if (!userMessage) return;
+
+        // Display user message
+        const userMessageElement = document.createElement('div');
+        userMessageElement.className = 'user-message';
+        userMessageElement.textContent = userMessage;
+        messages.appendChild(userMessageElement);
+
+        // Clear input
+        input.value = '';
+
+        // Send message to chatbot API
+        try {
+            const response = await fetch('/chatbot/query', {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/json'
+                },
+                body: JSON.stringify({ query: userMessage })
+            });
+            const data = await response.json();
+
+            // Display chatbot response
+            const botMessageElement = document.createElement('div');
+            botMessageElement.className = 'bot-message';
+            botMessageElement.textContent = data.response;
+            messages.appendChild(botMessageElement);
+        } catch (error) {
+            console.error('Error:', error);
+        }
+    }
+</script>
+
+<style>
+    .chatbot-container {
+        border: 1px solid #ccc;
+        border-radius: 5px;
+        padding: 10px;
+        max-width: 400px;
+        margin: 20px auto;
+    }
 
+    .chatbot-header {
+        text-align: center;
+        font-weight: bold;
+    }
+
+    .chatbot-messages {
+        height: 200px;
+        overflow-y: auto;
+        border: 1px solid #ccc;
+        padding: 10px;
+        margin-bottom: 10px;
+    }
 
+    .chatbot-input {
+        display: flex;
+    }
 
-{% endblock %}
+    .chatbot-input input {
+        flex: 1;
+        padding: 5px;
+    }
+
+    .chatbot-input button {
+        padding: 5px 10px;
+    }
+
+    .user-message {
+        text-align: right;
+        color: blue;
+    }
+
+    .bot-message {
+        text-align: left;
+        color: green;
+    }
+</style>
+
+{% endblock %}