File size: 5,451 Bytes
24773d4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
"""
Handbook Generator — Gradio UI (legacy fallback).
Primary UI is streamlit_app.py.
Run: python app.py
"""
import asyncio
import shutil
from pathlib import Path

import gradio as gr

from config import GROK_API_KEY, UPLOADS_DIR, BASE_DIR
from handbook_generator import build_handbook
from rag import get_context_for_query, index_pdf, reset_index

HANDBOOK_EXPORT_PATH = BASE_DIR / "handbook_export.md"


def _run_async(coro):
    """Run an async coroutine from sync Gradio code."""
    try:
        return asyncio.run(coro)
    except RuntimeError:
        import concurrent.futures
        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
            return pool.submit(asyncio.run, coro).result(timeout=300)


def ensure_api_key():
    if not GROK_API_KEY:
        raise gr.Error(
            "GROK_API_KEY is not set. Create a .env file in the ass2 folder with: GROK_API_KEY=your-key"
        )


def _file_path(f):
    """Get path from Gradio file input (path string or object with .name)."""
    if f is None:
        return None
    if isinstance(f, (str, Path)):
        return Path(f)
    return Path(getattr(f, "name", str(f)))


def upload_and_index(files):
    """Handle PDF upload(s) and index into LightRAG."""
    if not files:
        return "No files selected."
    ensure_api_key()
    reset_index()
    saved = []
    for f in (files if isinstance(files, list) else [files]):
        path = _file_path(f)
        if path is None or not path.exists():
            continue
        dest = UPLOADS_DIR / path.name
        try:
            shutil.copy(str(path), str(dest))
        except Exception:
            dest = path
        try:
            n = _run_async(index_pdf(dest, source_name=path.name))
            saved.append(f"{path.name}: indexed")
        except Exception as e:
            saved.append(f"{path.name}: Error - {e}")
    return "\n".join(saved) if saved else "No PDFs processed."


def chat(message, history):
    """RAG chat: retrieve context and answer using Grok via LiteLLM."""
    ensure_api_key()
    from litellm import completion
    from config import CHAT_MODEL

    context = _run_async(get_context_for_query(message))
    if not context or not context.strip():
        context = "No documents have been uploaded yet. Ask the user to upload PDFs first."

    system = (
        "You are a helpful assistant. Answer based ONLY on the following context "
        "from the user's uploaded documents. If the answer is not in the context, say so clearly."
    )
    user_content = f"Context from uploaded documents:\n\n{context}\n\n---\n\nUser question: {message}"

    resp = completion(
        model=CHAT_MODEL,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user_content},
        ],
        api_key=GROK_API_KEY,
        max_tokens=1500,
        temperature=0.3,
    )
    return (resp.choices[0].message.content or "").strip()


def run_handbook_simple(topic):
    """Generate handbook and return (status, markdown)."""
    ensure_api_key()
    if not (topic and topic.strip()):
        return "Enter a topic first.", ""
    status_msgs = []
    try:
        full_md = _run_async(build_handbook(topic.strip(), on_progress=status_msgs.append))
        status = "\n".join(status_msgs) if status_msgs else "Done."
        return status, full_md
    except Exception as e:
        return f"Error: {e}", ""


with gr.Blocks(title="Handbook Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Handbook Generator\nUpload PDFs, chat about them, and generate a 20,000+ word handbook.")

    with gr.Tab("Upload PDFs"):
        file_input = gr.File(
            file_count="multiple",
            file_types=[".pdf"],
            label="Upload one or more PDFs",
        )
        index_btn = gr.Button("Index PDFs")
        index_out = gr.Textbox(label="Index result", lines=4)

    with gr.Tab("Chat"):
        chatbot = gr.ChatInterface(
            fn=chat,
            type="messages",
            title="Ask questions about your uploaded documents",
        )

    with gr.Tab("Generate Handbook"):
        gr.Markdown(
            "Enter a topic (e.g. *Create a handbook on Retrieval-Augmented Generation*). "
            "Generation may take several minutes."
        )
        topic_in = gr.Textbox(
            label="Handbook topic",
            placeholder="e.g. Retrieval-Augmented Generation",
            lines=1,
        )
        gen_btn = gr.Button("Generate 20k-word handbook")
        status_out = gr.Textbox(label="Status", lines=4, interactive=False)
        handbook_out = gr.Markdown(label="Handbook (Markdown)")
        export_btn = gr.DownloadButton("Export as Markdown", visible=False)

    index_btn.click(
        fn=lambda files: upload_and_index(files) if files else "No files selected.",
        inputs=[file_input],
        outputs=[index_out],
    )

    def do_handbook(topic):
        status, md = run_handbook_simple(topic)
        if md:
            HANDBOOK_EXPORT_PATH.write_text(md, encoding="utf-8")
        return (
            status,
            md,
            gr.update(visible=bool(md), value=str(HANDBOOK_EXPORT_PATH) if md else None),
        )

    gen_btn.click(
        fn=do_handbook,
        inputs=[topic_in],
        outputs=[status_out, handbook_out, export_btn],
    )

if __name__ == "__main__":
    demo.launch(server_name="127.0.0.1", server_port=7860)