shovo896/codedocumentation

1import os2import sys3 4import gradio as gr5 6from ingest import ingest_repository, normalize_repo_url7from retriever import get_qa_chain8 9 10chain_cache = {}11 12 13def get_cached_chain(repo_info):14    namespace = repo_info["namespace"]15    if namespace not in chain_cache:16        chain_cache[namespace] = get_qa_chain(17            namespace=namespace,18            repo_url=repo_info["repo_url"],19            branch=repo_info["branch"],20        )21    return chain_cache[namespace]22 23 24def format_status(repo_info):25    return (26        f"Loaded repository: {repo_info['repo_url']}\n"27        f"Branch: {repo_info['branch']}\n"28        f"Files indexed: {repo_info['files']}\n"29        f"Chunks indexed: {repo_info['chunks']}"30    )31 32 33def load_repository(repo_url, branch):34    try:35        repo_url = normalize_repo_url(repo_url)36        branch = (branch or "").strip()37        repo_info = ingest_repository(repo_url, branch)38        chain_cache.pop(repo_info["namespace"], None)39        get_cached_chain(repo_info)40        return repo_info, format_status(repo_info), "", ""41    except Exception as e:42        print(f"Repository load error: {e}", file=sys.stderr)43        return None, f"Error: {e}", "", ""44 45 46def answer_query(repo_url, branch, question, repo_state):47    if not repo_url.strip():48        return repo_state, "Please enter a GitHub repository URL.", "", ""49 50    if not question.strip():51        return repo_state, "Please enter a question.", "", ""52 53    try:54        repo_url = normalize_repo_url(repo_url)55        branch = (branch or "").strip()56        state_matches = (57            repo_state58            and repo_state.get("repo_url") == repo_url59            and (not branch or repo_state.get("branch") == branch)60        )61 62        if not state_matches:63            repo_state = ingest_repository(repo_url, branch)64            chain_cache.pop(repo_state["namespace"], None)65 66        chain = get_cached_chain(repo_state)67        result = chain.invoke({"query": question})68        answer = result["result"]69 70        sources = set()71        for doc in result["source_documents"]:72            source = doc.metadata.get("source", "unknown")73            sources.add(source)74 75        source_text = "\n".join(sorted(sources))76        return repo_state, format_status(repo_state), answer, source_text77 78    except Exception as e:79        print(f"Query error: {e}", file=sys.stderr)80        return repo_state, f"Error: {e}", "", ""81 82 83with gr.Blocks(title="Code Doc Search") as demo:84    repo_state = gr.State(None)85 86    gr.Markdown("Code Documentation Search")87    gr.Markdown(88        "Paste any public GitHub repository URL, then ask questions about that codebase. "89        "Answers are always returned in English."90    )91 92    repo_url = gr.Textbox(93        label="GitHub Repository URL",94        placeholder="https://github.com/owner/repository",95        lines=1,96    )97 98    branch = gr.Textbox(99        label="Branch",100        value="",101        placeholder="Leave empty to use the repository default branch",102        lines=1,103    )104 105    load_btn = gr.Button("Load Repository", variant="secondary")106    status_box = gr.Textbox(label="Repository Status", lines=4, interactive=False)107 108    query = gr.Textbox(109        label="Your Question",110        placeholder="e.g. How does authentication work in this repository?",111        lines=2,112    )113 114    search_btn = gr.Button("Search", variant="primary")115 116    with gr.Row():117        answer_box = gr.Textbox(label="Answer", lines=10)118        source_box = gr.Textbox(label="Sources", lines=10)119 120    load_btn.click(121        fn=load_repository,122        inputs=[repo_url, branch],123        outputs=[repo_state, status_box, answer_box, source_box],124    )125 126    search_btn.click(127        fn=answer_query,128        inputs=[repo_url, branch, query, repo_state],129        outputs=[repo_state, status_box, answer_box, source_box],130    )131 132if __name__ == "__main__":133    server_port = int(os.environ.get("PORT", "7860"))134    demo.queue(default_concurrency_limit=1)135    demo.launch(136        theme=gr.themes.Soft(),137        server_name="0.0.0.0",138        server_port=server_port,139        share=False,140        prevent_thread_lock=False,141        debug=True,142        show_error=True,143    )144