app.py

import streamlit as st
from smolagents import Tool, CodeAgent, HfApiModel
from langchain.text_splitter import RecursiveCharacterTextSplitter, MarkdownTextSplitter
from langchain_community.retrievers import BM25Retriever
from langchain.docstore.document import Document
from datasets import load_dataset, concatenate_datasets

st.set_page_config(
    page_title="African History Search Engine",
    page_icon="🌍",
    layout="wide"
)

class RetrieverTool(Tool):
    name = "retriever"
    description = "Uses BM25 search to retrieve relevant African historical documentation"
    inputs = {
        "query": {
            "type": "string",
            "description": "The historical query in affirmative form rather than a question"
        }
    }
    output_type = "string"

    def __init__(self, docs, k1=1.5, b=0.75, **kwargs):
        super().__init__(**kwargs)
        self.retriever = BM25Retriever.from_documents(
            docs,
            k=12,
            k1=k1,
            b=b
        )
        self.docs = docs
        self.avg_doc_length = sum(len(doc.page_content.split()) for doc in docs) / len(docs)

    def forward(self, query: str) -> str:  # Matches exactly with inputs
        # Preprocess query
        query = self._preprocess_query(query)
        
        # Retrieve documents
        docs = self.retriever.get_relevant_documents(query)

        # Format response
        main_response = "Retrieved documents (ranked by relevance):\n\n"
        
        for i, doc in enumerate(docs, 1):
            doc_length = len(doc.page_content.split())
            length_factor = doc_length / self.avg_doc_length
            
            main_response += f"Document {i} (Length Factor: {length_factor:.2f})\n"
            main_response += f"{doc.page_content}\n\n"
            
            if doc.metadata:
                main_response += f"Metadata: {doc.metadata}\n"
            main_response += "---\n\n"

        return main_response

    def _preprocess_query(self, query: str) -> str:
        question_words = ["what", "when", "where", "who", "why", "how"]
        query_terms = query.lower().split()
        if query_terms[0] in question_words:
            query_terms = query_terms[1:]
        return " ".join(query_terms)

# Process documents
def prepare_docs(documents):
    text_splitter = MarkdownTextSplitter(
        chunk_size=1000,
        chunk_overlap=200
    )
    return text_splitter.split_documents(documents)

# Initialize agent
def create_rag_agent(processed_docs):
    retriever_tool = RetrieverTool(processed_docs)
    return CodeAgent(
        tools=[retriever_tool],
        model=HfApiModel(),
        verbose=True
    )

def format_search_results(results: str):
    """Format the search results into main content and sources sections"""
    if "### 📚 Sources:" in results:
        main_content, sources = results.split("### 📚 Sources:")
        
        # Create two columns with adjusted ratios
        col1, col2 = st.columns([3, 2])
        
        with col1:
            st.markdown("### 📖 Main Findings")
            st.markdown(main_content)
        
        with col2:
            st.markdown("### 📚 Sources")
            st.markdown(sources, unsafe_allow_html=True)
    else:
        st.markdown(results)

@st.cache_resource
def get_agent():
    """Single function to handle data loading, processing, and agent creation"""
    # Load dataset
    dataset = load_dataset("Svngoku/African-History-Extra-11-30-24")
    train_docs = dataset["train"]
    test_docs = dataset["test"]
    source_docs = concatenate_datasets([train_docs, test_docs])
    
    # Create documents
    documents = [
        Document(
            page_content=item['content'],
            metadata={
                "source": item['url'],
                "title": item['title'],
                "description": item['description'],
                "published_time": item['publishedTime']
            }
        )
        for item in source_docs
    ]
    
    # Process documents
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=500,
        add_start_index=True,
        strip_whitespace=True,
    )
    processed_docs = text_splitter.split_documents(documents)
    
    # Create and return agent
    retriever_tool = RetrieverTool(processed_docs)
    return CodeAgent(
        tools=[retriever_tool],
        model=HfApiModel(),
    )

# Streamlit UI
st.title("🌍 African History Search Engine")
st.markdown("""
This search engine uses advanced AI to help you explore African history. 
It provides detailed, sourced information from a curated database of historical documents.
""")

# Initialize agent
if 'agent' not in st.session_state:
    with st.spinner("Loading historical database..."):
        st.session_state.agent = get_agent()

# Search interface
search_query = st.text_input(
    "🔍 Search African History",
    placeholder="E.g., Tell me about the Kingdom of Kush",
    help="Enter any question about African history"
)

# Advanced search options
with st.expander("Advanced Search Options"):
    search_type = st.radio(
        "Search Type",
        ["General Query", "Specific Time Period", "Geographic Region"],
        help="Select the type of search you want to perform"
    )
    
    if search_type == "Specific Time Period":
        search_query = f"Focus on the time period: {search_query}"
    elif search_type == "Geographic Region":
        search_query = f"Focus on the region of: {search_query}"

# Search button
if st.button("Search", type="primary"):
    if search_query:
        with st.spinner("Searching historical records..."):
            try:
                results = st.session_state.agent.run(search_query)
                
                # Use the formatter to display results
                format_search_results(results)
                
                # Add methodology note
                st.markdown("---")
                st.info("""
                💡 **How to read the results:**
                - Main findings are summarized on the left
                - Source references are numbered [Source X]
                - Click on source details on the right to expand
                - Follow the links to read the original articles
                """)
                
            except Exception as e:
                st.error(f"An error occurred during the search: {e}")
    else:
        st.warning("Please enter a search query to begin.")

# Sidebar with additional information
with st.sidebar:
    st.markdown("### About This Search Engine")
    st.markdown("""
    This search engine specializes in African history, providing:
    - 📚 Detailed historical information
    - 🔍 Source verification
    - 🌍 Geographic context
    - ⏳ Historical timeline context
    """)
    
    st.markdown("### Data Sources")
    st.markdown("Our database includes information from various historical documents, "
               "academic papers, and verified historical records.")

# Footer
st.markdown("---")
st.caption("Powered by SmolAgents, RAG, and African History Dataset")