simplify the code

Author: liyin2015

README.md

@@ -35,57 +35,7 @@ poetry run streamlit run app.py
   - `DatabaseManager` in `src/data_pipeline.py` to manage the database.
   - `RAG` class in `src/rag.py` to manage the whole RAG lifecycle.
 
-<!-- CREATE Checklist -->
+- [ ] Conditional retrieval. Sometimes users just want to clarify a past conversation, no extra context needed.
 - [ ] Create an evaluation dataset  
 - [ ] Evaluate the RAG performance on the dataset  
 - [ ] Auto-optimize the RAG model
-<!-- ## Learn
-
-## Local Storage
-We use adalflow's root directory, which is at ~/.adalflow.
-- repos/repo_name/...
-- repos/repo_name_db/...
-
-- data_pipeline.py: From the main and local code test, you will know the process of download repo and chunk files, and embed the chunks.
-- rag.py: The main code of the RAG model. -->
-
-<!-- ### Command Line Interface
-
-Run the RAG system directly:
-```bash
-poetry run python rag.py
-```
-
-## Usage Examples
-
-1. **Demo Version (app.py)**
-   - Ask about Alice (software engineer)
-   - Ask about Bob (data scientist)
-   - Ask about the company cafeteria
-   - Test memory with follow-up questions
-
-2. **Repository Analysis (app_repo.py)**
-   - Enter your repository path
-   - Click "Load Repository"
-   - Ask questions about classes, functions, or code structure
-   - View implementation details in expandable sections
-
-## Security Note
-
-- Never commit your `.streamlit/secrets.toml` file
-- Add it to your `.gitignore`
-- Keep your API key secure
-
-## Example Queries
-
-- "What does the RAG class do?"
-- "Show me the implementation of the Memory class"
-- "How is data processing handled?"
-- "Explain the initialization process"
-
-## TODO
-
-- [ ] Add evaluation metrics
-- [ ] Improve the embedding model
-- [ ] Improve the text splitter and chunking
-- [ ] Improve the retriever -->

app.py

@@ -5,6 +5,7 @@
     extract_class_definition,
     extract_class_name_from_query,
 )
+from typing import List
 
 from config import DEFAULT_GITHUB_REPO
 
@@ -51,9 +52,22 @@ def init_rag(repo_path_or_url: str):
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.write(message["content"])
-        if "context" in message:
-            with st.expander(f"View source from {message.get('file_path', 'unknown')}"):
-                st.code(message["context"], language=message.get("language", "python"))
+        # if "context" in message:
+        #     with st.expander(f"View source from {message.get('file_path', 'unknown')}"):
+        #         st.code(message["context"], language=message.get("language", "python"))
+
+from adalflow.core.types import Document
+
+
+def form_context(context: List[Document]):
+    formatted_context = ""
+    for doc in context:
+        formatted_context += ""
+        f"file_path: {doc.meta_data.get('file_path', 'unknown')} \n"
+        f"language: {doc.meta_data.get('type', 'python')} \n"
+        f"content: {doc.text} \n"
+    return formatted_context
+
 
 if st.session_state.rag and (
     prompt := st.chat_input(
@@ -65,51 +79,29 @@ def init_rag(repo_path_or_url: str):
     with st.chat_message("user"):
         st.write(prompt)
 
-    class_name = extract_class_name_from_query(prompt)
+    # class_name = extract_class_name_from_query(prompt)
+    query = prompt
 
     with st.chat_message("assistant"):
         with st.spinner("Analyzing code..."):
             response, docs = st.session_state.rag(prompt)
 
             # Show relevant context first, then the explanation
             if docs and docs[0].documents:
-                # Try to find implementation code first
-                implementation_docs = [
-                    doc
-                    for doc in docs[0].documents
-                    if doc.meta_data.get("is_implementation", False)
-                ]
-
-                # Use implementation if found, otherwise use first document
-                doc = (
-                    implementation_docs[0]
-                    if implementation_docs
-                    else docs[0].documents[0]
-                )
-                context = doc.text
-                file_path = doc.meta_data.get("file_path", "unknown")
-                file_type = doc.meta_data.get("type", "python")
-
-                # If asking about a specific class, try to extract just that class definition
-                if class_name and file_type == "python":
-                    class_context = extract_class_definition(context, class_name)
-                    if class_context != context:  # Only use if we found the class
-                        context = class_context
-
-                with st.expander(f"View source from {file_path}"):
-                    st.code(context, language=file_type)
+                context = docs[0].documents
 
                 # Now show the explanation
-                st.write(response)
+                st.write(f"Rationale: {response.rationale}")
+                st.write(f"Answer: {response.answer}")
+
+                st.write(f"context: {form_context(context)}")
 
                 # Add to chat history
                 st.session_state.messages.append(
                     {
                         "role": "assistant",
                         "content": response,
                         "context": context,
-                        "file_path": file_path,
-                        "language": file_type,
                     }
                 )
             else:

config.py

@@ -1,58 +1,5 @@
 from adalflow import OpenAIClient
 
-# Prompt templates for different use cases
-prompts = {
-    "code_analysis": r"""
-You are a code analysis assistant that helps users understand code implementations.
-
-Your task is to analyze code and explain its functionality, focusing on:
-1. Implementation details and how the code works
-2. Class methods, their purposes, and interactions
-3. Key algorithms and data structures used
-4. Code patterns and architectural decisions
-
-When analyzing code:
-- Be concise and focus on the most important aspects
-- Explain the main purpose and key functionality first
-- Highlight critical methods and their roles
-- Keep explanations clear and to the point
-
-When asked about a specific class or function:
-1. Start with a one-sentence overview
-2. List the key methods and their purposes
-3. Explain the main functionality
-4. Keep the explanation focused and brief
-
-Previous conversation history is provided to maintain context of the discussion.
-Use the conversation history to provide more relevant and contextual answers about the code.
-
-Output JSON format:
-{
-    "answer": "Concise explanation of the code implementation",
-}""",
-    "general_qa": r"""
-You are a helpful assistant answering questions about provided documents.
-
-Your task is to:
-1. Answer questions based on the provided context
-2. Use conversation history to maintain coherent dialogue
-3. Be clear and concise in your responses
-4. Stay factual and only use information from the context
-
-When responding:
-- Start with a direct answer to the question
-- Provide relevant details from the context
-- Maintain a friendly, conversational tone
-- If information is not in the context, say so
-
-Previous conversation history is provided to maintain context of the discussion.
-Use the conversation history to provide more relevant and contextual answers.
-
-Output JSON format:
-{
-    "answer": "Clear and concise response based on the context",
-}""",
-}
 
 configs = {
     "embedder": {
@@ -77,8 +24,8 @@
     },
     "text_splitter": {
         "split_by": "word",
-        "chunk_size": 100,
-        "chunk_overlap": 20,
+        "chunk_size": 400,
+        "chunk_overlap": 100,
     },
 }