| import os |
| import argparse |
| import subprocess |
| import time |
| import requests |
|
|
| |
| API_PORT = 8000 |
| FRONTEND_PORT = 7860 |
| MAX_PAGES = 1000 |
|
|
| def run_scraper(): |
| """Run the web scraper to collect data.""" |
| from buffalo_rag.scraper.scraper import BuffaloScraper |
| |
| if os.path.exists("data/raw"): |
| num_pages = len([name for name in os.listdir("data/raw") if os.path.isfile(os.path.join("data/raw", name))]) |
| if num_pages > MAX_PAGES: |
| print(f"{num_pages} scraped data files found under data/raw. Skipping data scraper stage.") |
| else: |
| print("Starting web scraper...") |
| scraper = BuffaloScraper() |
| scraper.scrape(max_pages=MAX_PAGES) |
| print("Scraping completed!") |
|
|
| def build_embeddings(): |
| """Process documents and create embeddings.""" |
| from buffalo_rag.embeddings.chunker import DocumentChunker |
| |
| print("Creating document chunks and embeddings...") |
| chunker = DocumentChunker() |
| chunks = chunker.create_chunks() |
| chunker.create_embeddings(chunks) |
| print("Embeddings created!") |
|
|
| def start_api_server(): |
| """Start the FastAPI backend server.""" |
| print("Starting API server...") |
| subprocess.run(["uvicorn", "buffalo_rag.api.main:app", "--host", "0.0.0.0", "--port", str(API_PORT), "--reload"]) |
|
|
| def start_flask_frontend(): |
| """Start the Flask frontend.""" |
| print("Starting Flask frontend...") |
| subprocess.run(["flask", "run", "--host=0.0.0.0", f"--port={FRONTEND_PORT}"], |
| env={**os.environ, "FLASK_APP": "buffalo_rag/frontend/flask_app.py"}) |
|
|
| def wait_for_server(url, timeout=30, interval=1): |
| """Waits for a server at the given URL to be reachable.""" |
| start_time = time.time() |
| print(f"Waiting for server at {url} to be ready...") |
| while time.time() - start_time < timeout: |
| try: |
| response = requests.get(url, timeout=interval) |
| if response.status_code < 500: |
| print(f"Server at {url} is ready.") |
| return True |
| except requests.exceptions.RequestException: |
| pass |
| time.sleep(interval) |
| print(f"Timeout waiting for server at {url}.") |
| return False |
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="BuffaloRAG - AI Assistant for UB International Students") |
| parser.add_argument("--scrape", action="store_true", help="Run web scraper") |
| parser.add_argument("--build", action="store_true", help="Build embeddings") |
| parser.add_argument("--api", action="store_true", help="Run API server") |
| parser.add_argument("--frontend", action="store_true", help="Run Flask frontend") |
| parser.add_argument("--all", action="store_true", help="Run the complete pipeline") |
| parser.add_argument("--run", action="store_true", help="Run frontend & backend servers") |
| |
| args = parser.parse_args() |
| |
| if args.scrape or args.all: |
| run_scraper() |
| |
| if args.build or args.all: |
| build_embeddings() |
| |
| if args.api or args.all or args.run: |
| if args.all or args.run: |
| api_process = subprocess.Popen(["uvicorn", "buffalo_rag.api.main:app", "--host", "0.0.0.0", "--port", str(API_PORT), "--reload"]) |
| api_ready = wait_for_server(f'http://localhost:{API_PORT}/', timeout=60) |
|
|
| if api_ready: |
| time.sleep(3) |
| flask_process = subprocess.Popen(["flask", "run", "--host=0.0.0.0", f"--port={FRONTEND_PORT}"], |
| env={**os.environ, "FLASK_APP": "buffalo_rag/frontend/flask_app.py"}) |
| |
| import webbrowser |
| webbrowser.open(f'http://localhost:{FRONTEND_PORT}') |
| |
| input("Press Enter to stop the server and exit...\n") |
| if api_process.poll() is None: |
| api_process.terminate() |
| if flask_process.poll() is None: |
| flask_process.terminate() |
|
|
| if api_process.poll() is None: |
| api_process.terminate() |
| else: |
| start_api_server() |
| |
| if args.frontend: |
| start_flask_frontend() |
|
|
| if __name__ == "__main__": |
| main() |
