Spaces:
Sleeping
Sleeping
Islam Mamedov commited on
Commit ·
3df6802
1
Parent(s): c94d0fe
Day 5 ablations: hybrid recall@20=0.97, bge-reranker-base hurts MRR (0.71->0.63)
Browse files- data/eval_cache.json +78 -1
- src/ask.py +15 -34
- src/eval.py +48 -58
- src/retrieval.py +126 -0
data/eval_cache.json
CHANGED
|
@@ -112,5 +112,82 @@
|
|
| 112 |
"9b0871cd909b0135": "I couldn't find this in the indexed codebase.",
|
| 113 |
"5f4a45bd0a897d76": "I couldn't find this in the indexed codebase.",
|
| 114 |
"95b2ceba4706cca8": "I couldn't find this in the indexed codebase.",
|
| 115 |
-
"615a9cbcb33f39c5": "I couldn't find this in the indexed codebase."
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 116 |
}
|
|
|
|
| 112 |
"9b0871cd909b0135": "I couldn't find this in the indexed codebase.",
|
| 113 |
"5f4a45bd0a897d76": "I couldn't find this in the indexed codebase.",
|
| 114 |
"95b2ceba4706cca8": "I couldn't find this in the indexed codebase.",
|
| 115 |
+
"615a9cbcb33f39c5": "I couldn't find this in the indexed codebase.",
|
| 116 |
+
"28bebeb9258a2ad9": "To set a custom status code for an endpoint in FastAPI, you can declare a parameter of type `Response` in your path operation function and set the `status_code` in that temporary response object. Here's an example [5]:\n```python\nfrom fastapi import FastAPI, Response\n\napp = FastAPI()\n\n@app.get(\"/\")\nasync def endpoint():\n return Response(content=\"Custom response\", status_code=201)\n```\nAlternatively, you can use the `status_code` parameter when defining the route [2]:\n```python\n@app.post(\"/\", status_code=201)\nasync def endpoint():\n return {\"message\": \"Success\"}\n```\nIf you need to dynamically set the status code, you can use a dependency to check the API key and raise an `HTTPException` with a custom status code [1]:\n```python\nfrom fastapi import Depends, HTTPException\nfrom fastapi.security import APIKeyHeader\n\nX_API_KEY = APIKeyHeader(name='X-API-Key')\n\ndef check_authentication_header(x_api_key: str = Depends(X_API_KEY)):\n if x_api_key != \"1234567890\":\n raise HTTPException(status_code=401, detail=\"Invalid API Key\")\n return {\"id\": 1234567890}\n```",
|
| 117 |
+
"025e758768958c60": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 118 |
+
"a1e4644f5631e3a4": "To receive an uploaded file in an endpoint, you can define a parameter of type `UploadFile` in your endpoint function. Here's an example [4]:\n```python\nfrom fastapi import FastAPI, UploadFile, File\n\napp = FastAPI()\n\n@app.post(\"/upload/\")\nasync def create_upload_file(file: UploadFile = File(...)):\n return {\"filename\": file.filename}\n```\nMake sure to install `python-multipart` first, as uploaded files are sent as \"form data\" [4]:\n```console\n$ pip install python-multipart\n```\nYou can also define multiple files and form fields at the same time using `File` and `Form` [5]:\n```python\nfrom fastapi import FastAPI, File, Form, UploadFile\n\napp = FastAPI()\n\n@app.post(\"/upload/\")\nasync def create_upload_file(\n file: UploadFile = File(...), \n form_data: str = Form(...)\n):\n return {\"filename\": file.filename, \"form_data\": form_data}\n```",
|
| 119 |
+
"c93032d71d708e33": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 120 |
+
"34b74e6aa3aed093": "You can define background tasks to be run after returning a response using the `BackgroundTasks` class from FastAPI. This is useful for operations that need to happen after a request, but that the client doesn't really have to be waiting for the operation to complete before receiving the response [1].\n\nHere's an example of how to use it:\n```python\nfrom fastapi import BackgroundTasks, FastAPI\n\napp = FastAPI()\n\ndef write_notification(email: str, message=\"\"):\n with open(\"log.txt\", mode=\"w\") as email_file:\n content = f\"notification for {email}: {message}\"\n email_file.write(content)\n\n@app.post(\"/send-notification/{email}\")\nasync def send_notification(email: str, background_tasks: BackgroundTasks):\n background_tasks.add_task(write_notification, email, message=\"some notification\")\n return {\"message\": \"Notification sent in the background\"}\n```\nIn this example, the `write_notification` function is run in the background after the response is sent [2].\n\nYou can also declare a parameter in a path operation function or dependency function with the type `BackgroundTasks`, and then use it to schedule the execution of background tasks after the response is sent [3].\n\nNote that if you return a `Response` object that already has its own `background` attribute set, the injected background tasks will be silently discarded. To avoid this, you can add a warning by emitting a `UserWarning` [5].",
|
| 121 |
+
"f295cb3f49e9f9eb": "```json\n{\n \"faithful\": false,\n \"correct\": true\n}\n```",
|
| 122 |
+
"dd0cb2c9debbdef4": "To enable CORS in your FastAPI application, you can use the `CORSMiddleware`. Here's how you can do it:\n\n* Import `CORSMiddleware`.\n* Create a list of allowed origins (as strings).\n* Add it as a \"middleware\" to your FastAPI application.\n\nYou can configure the `CORSMiddleware` with the following parameters:\n* `allow_origins`: A list of origins that should be permitted to make cross-origin requests.\n* `allow_methods`: A list of HTTP methods that should be allowed for cross-origin requests. Defaults to `['GET']`.\n* `allow_headers`: A list of HTTP request headers that should be supported for cross-origin requests. Defaults to `[]`.\n* `allow_credentials`: Indicate that cookies should be supported for cross-origin requests. Defaults to `False`.\n\nHere's an example:\n```python\nfrom fastapi import FastAPI\nfrom fastapi.middleware.cors import CORSMiddleware\n\napp = FastAPI()\n\norigins = [\n \"http://localhost:8000\",\n \"http://localhost:8001\",\n]\n\napp.add_middleware(\n CORSMiddleware,\n allow_origins=origins,\n allow_credentials=True,\n allow_methods=[\"*\"],\n allow_headers=[\"*\"],\n)\n```\nThis way, you can enable CORS for your frontend to call the API [1].",
|
| 123 |
+
"3a005c0235f0e466": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 124 |
+
"839285ca94ecb7e6": "You can declare optional query parameters with default values by setting their default to a specific value, such as `None`, or other default values. \n\nFor example, to declare an optional query parameter `q` with a default value of `None`, you can use:\n```python\nq: str | None = None\n```\nAlternatively, you can use `Query` to declare it explicitly as a query parameter:\n```python\nq: str | None = Query(default=None)\n```\nYou can also use other default values, such as strings, like `\"fixedquery\"`, and add validations, like `min_length`:\n```python\nq: str = \"fixedquery\"\n```\nOr \n```python\nq: Annotated[str, Query(min_length=3)] = \"fixedquery\"\n``` \n [2][1][3]",
|
| 125 |
+
"f842416650139f03": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 126 |
+
"2cd8377170776646": "You can declare the type of a path parameter in the function using standard Python type annotations. For example, to declare a path parameter `item_id` as an `int`, you can do [1]:\n```python\n# item_id is declared to be an int\n```\nThis will give you editor support inside of your function, with error checks, completion, etc. \n\nHere's an example:\n```python\nfrom fastapi import FastAPI\n\napp = FastAPI()\n\n@app.get(\"/items/{item_id}\")\nasync def read_item(item_id: int):\n return {\"item_id\": item_id}\n```",
|
| 127 |
+
"2d4a0180a291cd61": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 128 |
+
"b83a0f2b3b85924d": "To receive a JSON request body using a Pydantic model, you first need to import `BaseModel` from `pydantic`. Then, define your Pydantic model by inheriting from `BaseModel` and including the fields you expect in the JSON request body. Here's a basic example:\n\n```python\nfrom pydantic import BaseModel\n\nclass Item(BaseModel):\n name: str\n description: str\n price: float\n tax: float\n```\n\nYou can then use this `Item` model as a parameter in your route functions. FastAPI will automatically validate the incoming JSON request body against your Pydantic model [3].\n\nFor instance:\n\n```python\nfrom fastapi import FastAPI\nfrom pydantic import BaseModel\n\napp = FastAPI()\n\nclass Item(BaseModel):\n name: str\n description: str\n price: float\n tax: float\n\n@app.post(\"/items/\")\nasync def create_item(item: Item):\n return item\n```\n\nIn this example, FastAPI will expect a JSON request body like:\n\n```json\n{\n \"name\": \"Foo\",\n \"description\": \"The pretender\",\n \"price\": 42.0,\n \"tax\": 3.2\n}\n```\n\nAnd it will validate it against the `Item` model. If the JSON is valid, it will be available as the `item` parameter in your route function. If not, FastAPI will return a validation error. \n\nYou can also declare a body as a `dict` with specific key and value types, which allows receiving keys that you don't already know beforehand [1]. \n\nAdditionally, Pydantic supports more complex types and validation, including nested models, lists, and optional fields. For example:\n\n```python\nfrom pydantic import BaseModel\nfrom typing import List, Optional\n\nclass Detail(BaseModel):\n some_data: Optional[List[str]]\n\nclass Item(BaseModel):\n name: str\n description: str\n price: float\n tax: float\n details: Optional[Detail]\n```",
|
| 129 |
+
"b055668cea673be6": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 130 |
+
"912194ce87a09163": "To receive form data instead of JSON, you can use `Form` in your FastAPI application. However, to use forms, you first need to install `python-multipart`. You can install it using pip:\n\n```console\n$ pip install python-multipart\n```\n\nThen, you can declare form fields using `Form` parameters in your path operations. Note that when using `Form` parameters, you can't also declare `Body` fields that you expect to receive as JSON, as the request will have the body encoded using `application/x-www-form-urlencoded` instead of `application/json` [1][2].\n\nHere's an example of how to define form parameters:\n\n```python\nfrom fastapi import FastAPI, Form\n\napp = FastAPI()\n\n@app.post(\"/login/\")\ndef login(username: str = Form(...), password: str = Form(...)):\n return {\"username\": username}\n```\n\nIn this example, `username` and `password` are form fields that will be received from the request [1].",
|
| 131 |
+
"0650db67e0d2f9db": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 132 |
+
"fc10baea436a99b2": "To return a custom error response with a specific status code and message in FastAPI, you can create a custom `Response` object. For example, you can use `JSONResponse` to return a JSON response with a specific status code and message.\n\nHere's an example using `JSONResponse` directly [1]:\n```python\nfrom fastapi import FastAPI\nfrom fastapi.responses import JSONResponse\n\napp = FastAPI()\n\n@app.get(\"/items/\")\nasync def read_items():\n return JSONResponse(status_code=422, content={\"detail\": \"Custom error message\"})\n```\nAlternatively, you can use exception handlers to catch specific exceptions and return custom error responses. For example, you can use the `request_validation_exception_handler` in `fastapi/exception_handlers.py` as a reference [5]:\n```python\nfrom fastapi import FastAPI, Request\nfrom fastapi.exceptions import RequestValidationError\nfrom fastapi.responses import JSONResponse\nfrom fastapi import status\n\napp = FastAPI()\n\n@app.exception_handler(RequestValidationError)\nasync def validation_exception_handler(request: Request, exc: RequestValidationError):\n return JSONResponse(\n status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,\n content={\"detail\": \"Custom validation error message\"},\n )\n```\nYou can also use the `response_class` parameter in your path operation decorator to specify a custom response class, and then return a custom response object [3]:\n```python\nfrom fastapi import FastAPI\nfrom fastapi.responses import HTMLResponse\n\napp = FastAPI()\n\n@app.get(\"/items/\", response_class=HTMLResponse)\nasync def read_items():\n return HTMLResponse(content=\"<html>Custom error message</html>\", status_code=422)\n```",
|
| 133 |
+
"4be371a4509f7b0f": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 134 |
+
"31258a5bd2aabad3": "You can use the `response_model` parameter in FastAPI to control which fields are returned in the response. This is particularly useful when you want to filter data to expose only public fields and avoid exposing sensitive information.\n\nHere's an example:\n```python\nfrom fastapi import FastAPI\nfrom pydantic import BaseModel\n\napp = FastAPI()\n\nclass InternalItem(BaseModel):\n name: str\n description: str | None = None\n secret_key: str\n\nclass Item(BaseModel):\n name: str\n description: str | None = None\n\n@app.get(\"/items/me\", response_model=Item)\nasync def get_item() -> Any:\n item = InternalItem(\n name=\"Foo\", description=\"A very nice Item\", secret_key=\"supersecret\"\n )\n return item\n```\nIn this example, even though the `get_item` function returns an `InternalItem` object, which has a `secret_key` attribute, the `response_model` parameter is set to `Item`, which does not include the `secret_key` field. As a result, the JSON response sent to the client will not include the `secret_key` field [2].\n\nYou can use `response_model` to define response models and ensure private data is filtered out. Additionally, you can use `response_model_exclude_unset` to return only the values explicitly set [3].",
|
| 135 |
+
"90b935eea78b3d8e": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 136 |
+
"747a7a8dd0005c6b": "You can serve static files like images or CSS using the `StaticFiles` class in FastAPI. To do this, you need to import `StaticFiles` from `fastapi.staticfiles` and then mount a `StaticFiles()` instance in a specific path [1].\n\nHere's an example of how to serve static files:\n```python\nfrom fastapi.staticfiles import StaticFiles\n\n# Mount the StaticFiles instance\napp.mount(\"/static\", StaticFiles(directory=\"static\"), name=\"static\")\n```\nIn this example, the `StaticFiles` instance is mounted at the `/static` path and serves files from the `static` directory.\n\nYou can then access your static files by visiting `http://127.0.0.1:8000/static/filename.css` (or `.js`, `.jpg`, etc.) [4].\n\nFor example, if you have a CSS file `styles.css` in the `static` directory, you can link to it in your HTML template using `url_for()`:\n```jinja\n<link rel=\"stylesheet\" href=\"{{ url_for('static', path='styles.css') }}\">\n```\nThis way, FastAPI will serve the `styles.css` file automatically at the URL `/static/styles.css` [3].",
|
| 137 |
+
"5ca0d3bce8d33fdc": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 138 |
+
"4505811c313090f9": "To write tests for your FastAPI endpoints, you can use the `TestClient` provided by FastAPI. Here's a step-by-step guide:\n\n1. Install `httpx` using pip: `pip install httpx`.\n2. Import `TestClient` and create a `TestClient` instance by passing your FastAPI application to it.\n3. Write test functions with names starting with `test_`.\n4. Use the `TestClient` object to send requests to your FastAPI application.\n5. Write simple `assert` statements to check the responses.\n\nExample:\n```python\nfrom fastapi.testclient import TestClient\nfrom myapp import app\n\nclient = TestClient(app)\n\ndef test_read_main():\n response = client.get(\"/\")\n assert response.status_code == 200\n assert response.json() == {\"message\": \"Hello World\"}\n```\nNote that you can use `pytest` directly with FastAPI.\n\nIf you need to test asynchronous functions, you can use `async def` test functions and `await` the `TestClient` methods [2].\n\nFor more information, check out the [Testing](docs/en/docs/tutorial/testing.md) and [Async Tests](docs/en/docs/advanced/async-tests.md) sections in the FastAPI documentation [1][4].",
|
| 139 |
+
"300e299e5af00fb6": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 140 |
+
"18929192fee244d9": "To use WebSockets in FastAPI, you first need to create a WebSocket endpoint using the `@app.websocket()` decorator. Here's a basic example [5]:\n```python\nfrom fastapi import FastAPI, WebSocket\n\napp = FastAPI()\n\n@app.websocket(\"/ws\")\nasync def websocket_endpoint(websocket: WebSocket):\n await websocket.accept()\n while True:\n data = await websocket.receive_text()\n await websocket.send_text(f\"Message text was: {data}\")\n```\nThis code sets up a WebSocket endpoint at `/ws` that accepts WebSocket connections, receives text messages, and sends responses back to the client.\n\nYou can run your FastAPI application using `fastapi dev` and access it at `http://127.0.0.1:8000` [1].\n\nFor more information on handling WebSocket connections, including disconnections and multiple clients, refer to the FastAPI documentation [3][4].\n\nThe `WebSocket` class is provided directly by Starlette, but you can import it from `fastapi` for convenience [2]:\n```python\nfrom fastapi import WebSocket\n```",
|
| 141 |
+
"f6d46715b83ee69b": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 142 |
+
"0906e322d446b450": "You can set a cookie in a response by declaring a parameter of type `Response` in your path operation function and then setting the cookie in that temporary response object [1].\n\nHere's an example of how to set a cookie:\n```python\nfrom fastapi import FastAPI, Response\n\napp = FastAPI()\n\n@app.get(\"/\")\ndef read_root(response: Response):\n response.set_cookie(\"cookie_name\", \"cookie_value\")\n return {\"message\": \"Hello World\"}\n```\nAlternatively, you can also create a response directly and set cookies in it [2]:\n```python\nfrom fastapi import FastAPI, JSONResponse\n\napp = FastAPI()\n\n@app.get(\"/\")\ndef read_root():\n response = JSONResponse(content={\"message\": \"Hello World\"}, media_type=\"application/json\")\n response.set_cookie(\"cookie_name\", \"cookie_value\")\n return response\n```\nYou can import the `Response` class directly from `fastapi` [3]:\n```python\nfrom fastapi import Response\n```",
|
| 143 |
+
"43e79d44a1cf2a87": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 144 |
+
"33b387b4db40638d": "You can add custom headers to a response in several ways. \n\n1. **Using the `X-` prefix**: You can add custom proprietary headers using the `X-` prefix. However, if you want a client in a browser to see these headers, you need to add them to your CORS configurations using the `expose_headers` parameter [1][2].\n\n2. **Return a `Response` directly**: You can create a response and pass the headers as an additional parameter [3].\n\n```python\nfrom fastapi import FastAPI, Response\n\napp = FastAPI()\n\n@app.get(\"/\")\ndef read_root():\n headers = {\"X-Custom-Header\": \"Custom Value\"}\n return Response(content=\"Hello, World!\", media_type=\"text/plain\", headers=headers)\n```\n\n3. **Use a `Response` parameter**: You can declare a parameter of type `Response` in your path operation function and set headers in that temporary response object [5].\n\n```python\nfrom fastapi import FastAPI, Response\n\napp = FastAPI()\n\n@app.get(\"/\")\ndef read_root(response: Response):\n response.headers[\"X-Custom-Header\"] = \"Custom Value\"\n return {\"message\": \"Hello, World!\"}\n```\n\n4. **Middleware**: You can also add custom headers using middleware. You can create a middleware function that receives the request and a function `call_next` that will receive the request as a parameter, and then modify the response before returning it [2].\n\n```python\nfrom fastapi import FastAPI, Request\n\napp = FastAPI()\n\n@app.middleware(\"http\")\nasync def add_custom_header(request: Request, call_next):\n response = await call_next(request)\n response.headers[\"X-Custom-Header\"] = \"Custom Value\"\n return response\n```",
|
| 145 |
+
"e7df91dc36af9185": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 146 |
+
"d336527362c38e67": "To implement JWT token authentication, you'll need to handle JWT tokens and update the `/token` path operation. \n\nFirst, import the necessary modules, create a random secret key, and define the algorithm and expiration time for the token. Then, define a Pydantic model for the token endpoint response and create a utility function to generate a new access token [2].\n\nNext, update the `/token` path operation to create a `timedelta` with the expiration time and return a real JWT access token. Use the `sub` key to store the user's identification, ensuring it's unique and a string [3].\n\nHere's a high-level example of the steps:\n```python\nimport secrets\nimport jwt\nfrom datetime import datetime, timedelta\n\n# Define the secret key, algorithm, and expiration time\nSECRET_KEY = secrets.token_urlsafe(32)\nALGORITHM = \"HS256\"\nACCESS_TOKEN_EXPIRE_MINUTES = 30\n\n# Define the Pydantic model for the token endpoint response\nclass Token(BaseModel):\n access_token: str\n token_type: str\n\n# Create a utility function to generate a new access token\ndef create_access_token(data: dict, expires_delta: timedelta | None = None):\n to_encode = data.copy()\n if expires_delta:\n expire = datetime.utcnow() + expires_delta\n else:\n expire = datetime.utcnow() + timedelta(minutes=15)\n to_encode.update({\"exp\": expire})\n encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)\n return encoded_jwt\n\n# Update the /token path operation to return a JWT access token\n@app.post(\"/token\", response_model=Token)\nasync def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends()):\n # Authenticate the user and return a JWT access token\n access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)\n access_token = create_access_token(\n data={\"sub\": form_data.username}, expires_delta=access_token_expires\n )\n return {\"access_token\": access_token, \"token_type\": \"bearer\"}\n```",
|
| 147 |
+
"fb49c2bfc14f427f": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 148 |
+
"4b5e793ae7400937": "FastAPI's dependency injection system allows you to declare dependencies for your path operation functions, which are then executed and injected with the results by FastAPI. These dependencies are functions that can take the same parameters as path operation functions and return any value. You declare a dependency using the `Depends` function, passing the dependency function as an argument, and FastAPI will take care of executing it and injecting the results [1][3].\n\nDependencies can be used for various purposes, such as [2]:\n* Handling logic that can't be declared in Pydantic validation and requires additional logic\n* Managing external resources or blocking operations\n* Creating sub-dependencies\n* Sharing logic across multiple endpoints\n* Handling cleanup with `yield`, such as database sessions or file handles\n\nYou can apply shared dependencies at the router level using `dependencies=[Depends(...)]`. The dependency injection system also supports more advanced patterns, including `yield` with `scope` and class dependencies [4].\n\nThe dependency injection system integrates well with other FastAPI features, such as background tasks, allowing you to declare `BackgroundTasks` parameters in path operation functions, dependencies, or sub-dependencies [5].",
|
| 149 |
+
"fa4678ce5f5b85e9": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 150 |
+
"6d5309decb05a061": "`jsonable_encoder` is a function in FastAPI that converts an object, such as a Pydantic model, to a JSON-compatible version. This is necessary when you need to store or return data that isn't natively compatible with JSON, such as `datetime` objects or NumPy types.\n\nIt takes an object and returns a Python standard data structure (e.g., a `dict`) with values and sub-values that are all compatible with JSON. For example, it can convert a `datetime` object to a string in ISO format and a Pydantic model to a `dict`.\n\nYou need to use `jsonable_encoder` in scenarios like [1][2][4]:\n* Storing data in a database that only accepts JSON-compatible data.\n* Returning Pydantic models or other non-JSON-compatible data directly in a response.\n* Handling data with non-JSON-compatible types, such as NumPy types from Pandas DataFrames.\n\nHere's an example of using `jsonable_encoder`:\n```python\nfrom fastapi.encoders import jsonable_encoder\nfrom datetime import datetime\n\nclass Item:\n def __init__(self, name: str, created_at: datetime):\n self.name = name\n self.created_at = created_at\n\nitem = Item(\"Example Item\", datetime.now())\njson_compatible_item = jsonable_encoder(item)\n```\nIn this case, `jsonable_encoder` would convert the `item` object to a `dict` with a JSON-compatible representation of the `datetime` object. \n\nFastAPI internally uses `jsonable_encoder` to convert data, but it's also useful in other scenarios [1]. \n\nYou can also use it to handle Pandas DataFrames, which can contain non-JSON-compatible NumPy types [5]:\n```python\nimport pandas as pd\nfrom fastapi.encoders import jsonable_encoder\n\ndf = pd.DataFrame({\n \"count\": [1, 2, 3],\n \"score\": [1.5, float(\"nan\"), 3.0],\n})\n\njson_compatible_df = jsonable_encoder(df.to_dict(orient=\"records\"))\n```",
|
| 151 |
+
"6fb559b5ce1e6285": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 152 |
+
"369742ae24b4bc16": "To address your question, I need to understand what you're asking. You haven't specified a particular question about the FastAPI codebase, but I can try to infer and provide information based on the context provided.\n\nIf you're asking about the difference between `async def` and `def` endpoints in FastAPI, here [5] is some relevant documentation:\n\nIn FastAPI, you can define path operation functions using either `async def` or `def`. The main difference lies in how they handle asynchronous operations.\n\n- Use `async def` for functions that are asynchronous and need to use `await` with third-party libraries that support asynchronous operations.\n- Use `def` for synchronous functions or when working with third-party libraries that do not support asynchronous operations. FastAPI will run these functions in a thread pool.\n\nHere's a brief example:\n\n```python\nfrom fastapi import FastAPI\n\napp = FastAPI()\n\n# Asynchronous function\n@app.get('/async')\nasync def read_async():\n # Assume some_async_library is an async library\n data = await some_async_library.fetch_data()\n return data\n\n# Synchronous function\n@app.get('/sync')\ndef read_sync():\n # Assume some_sync_library is a sync library\n data = some_sync_library.fetch_data()\n return data\n```\n\nFor more detailed information, you can refer to the [documentation](docs/en/docs/async.md) and [skills documentation](fastapi/.agents/skills/fastapi/SKILL.md) provided.\n\nIf your question was something else, please provide more details. \n\nI couldn't find this in the indexed codebase for anything else.",
|
| 153 |
+
"f8e01cc97294352b": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 154 |
+
"7c2093facbe24099": "To run cleanup code after the response is sent, a dependency can use `yield` with the default scope `\"request\"`. This ensures that the exit code after `yield` is executed after the response is sent back to the client [1].\n\nHere's an example:\n```python\nfrom typing import Annotated\n\nfrom fastapi import Depends, FastAPI\n\napp = FastAPI()\n\ndef get_db():\n db = DBSession()\n try:\n yield db\n finally:\n db.close()\n\nDBDep = Annotated[DBSession, Depends(get_db)]\n\n@app.get(\"/items/\")\nasync def read_items(db: DBDep):\n return db.query(Item).all()\n```\nIn this case, the `db.close()` statement will be executed after the response is sent [1].\n\nAlternatively, you can explicitly specify the scope as `\"request\"`:\n```python\nUserDep = Annotated[DBSession, Depends(get_db, scope=\"request\")]\n```\nThis achieves the same result [1][4].",
|
| 155 |
+
"37abbb385dc1c1fc": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 156 |
+
"2ea7cf2fa8800910": "You can customize or extend the generated OpenAPI schema in several ways.\n\n## Using `openapi_extra`\n\nYou can extend the OpenAPI schema for a path operation using the `openapi_extra` parameter. This is useful for declaring OpenAPI extensions [1].\n\n## Modifying the OpenAPI Schema\n\nYou can also generate the OpenAPI schema and override each part that you need using the utility function `get_openapi`. This allows you to add custom extensions, such as a custom logo for ReDoc [4].\n\n## Extra JSON Schema Data in Pydantic Models\n\nAdditionally, you can declare extra JSON schema data in Pydantic models using the `model_config` attribute and setting `\"json_schema_extra\"` with a `dict` containing any additional data [3].\n\n## Overriding the Defaults\n\nTo extend OpenAPI, you can override the defaults by generating the OpenAPI schema, modifying it, and caching it. This involves creating a custom `openapi()` function and replacing the original one [4].\n\n## Example Code\n\nHere is an example of using `openapi_extra`:\n```python\nfrom fastapi import FastAPI\nfrom fastapi.openapi.models import OpenAPI\n\napp = FastAPI()\n\n@app.get(\"/items/\", openapi_extra={\n \"x-aperture-labs-portal\": \"stargate\"\n})\nasync def read_items():\n return [{\"item\": \"Item 1\"}]\n```\nAnd here is an example of generating and modifying the OpenAPI schema:\n```python\nfrom fastapi import FastAPI\nfrom fastapi.openapi.utils import get_openapi\n\napp = FastAPI()\n\ndef custom_openapi():\n if app.openapi_schema:\n return app.openapi_schema\n openapi_schema = get_openapi(\n title=app.title,\n version=app.version,\n description=app.description,\n routes=app.routes,\n )\n openapi_schema[\"info\"][\"x-logo\"] = {\"url\": \"https://example.com/logo.png\"}\n app.openapi_schema = openapi_schema\n return openapi_schema\n\napp.openapi = custom_openapi\n```",
|
| 157 |
+
"1c7b5a769de3e356": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 158 |
+
"9461901effdb2eba": "When request validation fails in FastAPI, it returns a `ValidationError` with a 422 status code. The error response includes details about the validation errors, such as the location and message of the error.\n\nThe default error response for request validation errors in FastAPI is:\n```json\n{\n \"detail\": [\n {\n \"loc\": [\n \"path\",\n \"item_id\"\n ],\n \"msg\": \"value is not a valid integer\",\n \"type\": \"type_error.integer\"\n }\n ]\n}\n```\nTo customize this error response, you can override the default exception handler for `RequestValidationError` using the `@app.exception_handler(RequestValidationError)` decorator [3].\n\nFor example, you can return a plain text response instead of JSON:\n```python\nfrom fastapi import FastAPI, RequestValidationError\nfrom fastapi.responses import PlainTextResponse\n\napp = FastAPI()\n\n@app.exception_handler(RequestValidationError)\nasync def validation_error_handler(request, exc):\n return PlainTextResponse(\n content=\"Validation errors: \" + str(exc),\n status_code=422,\n )\n```\nAlternatively, you can modify the OpenAPI entry for `ValidationError` and `HTTPValidationError` using the `openapi_schema` attribute of the FastAPI app [4].\n\nFastAPI uses Pydantic models to validate request data, and the `ValidationError` exception is raised when the data is invalid [1]. You can use Pydantic models to define the expected structure and content of request bodies.\n\nFastAPI also provides a `jsonable_encoder` function to encode the error details in a JSON-compatible format [5].",
|
| 159 |
+
"f8aadfea4bd68d7b": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 160 |
+
"4d260caf05bd6b11": "The OAuth2 password flow in FastAPI works as follows:\n\n* The user types their `username` and `password` in the frontend and hits `Enter`.\n* The frontend sends the `username` and `password` to a specific URL in the API (declared with `tokenUrl=\"token\"`).\n* The API checks the `username` and `password` and responds with a \"token\".\n* The frontend stores the token temporarily.\n* When the frontend needs to access an authenticated endpoint, it sends an `Authorization` header with the token.\n\nIn FastAPI, this flow is implemented using the `OAuth2PasswordRequestForm` and `OAuth2PasswordBearer` classes. \n\nThe `OAuth2PasswordRequestForm` class is used to collect the `username` and `password` as form data, while the `OAuth2PasswordBearer` class is used to handle the token authentication.\n\nHere's an example of how to use these classes:\n```python\nfrom fastapi import Depends, FastAPI\nfrom fastapi.security import OAuth2PasswordRequestForm, OAuth2PasswordBearer\n\napp = FastAPI()\n\noauth2_scheme = OAuth2PasswordBearer(tokenUrl=\"token\")\n\n@app.post(\"/login\")\ndef login(form_data: OAuth2PasswordRequestForm = Depends()):\n # Check username and password\n # Return token\n pass\n\n@app.get(\"/protected\")\ndef protected(token: str = Depends(oauth2_scheme)):\n # Use token to authenticate\n pass\n```\nThe `OAuth2PasswordRequestForm` class ensures that the `username` and `password` are collected using form data, while the `OAuth2PasswordBearer` class handles the token authentication [2][3].",
|
| 161 |
+
"a46ee8fe3b121777": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 162 |
+
"5dee9d2d67fcfdd8": "You can run code on application startup and shutdown using event handlers or the `lifespan` parameter.\n\n## Using Event Handlers\n\nYou can define event handlers using `@app.on_event(\"startup\")` and `@app.on_event(\"shutdown\")` decorators. Here's an example [1]:\n```python\n@app.on_event(\"startup\")\nasync def startup_event():\n # Code to run on startup\n pass\n\n@app.on_event(\"shutdown\")\nasync def shutdown_event():\n # Code to run on shutdown\n pass\n```\n## Using Lifespan Parameter\n\nAlternatively, you can use the `lifespan` parameter when creating the FastAPI app [2]:\n```python\nfrom contextlib import asynccontextmanager\n\n@asynccontextmanager\nasync def lifespan(app: FastAPI):\n # Code to run on startup\n try:\n yield\n finally:\n # Code to run on shutdown\n pass\n\napp = FastAPI(lifespan=lifespan)\n```\nIn both cases, you can execute code before the application starts receiving requests and after it finishes handling requests [3].",
|
| 163 |
+
"b8457d6b5cb39455": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 164 |
+
"01c7e4c1799ac5e9": "Yes, you can mount another application under a path prefix. This is achieved by using the `mount` method on a FastAPI instance. \n\nFor example, if you have a sub-application `subapi` and you want to mount it under the path prefix `/api/v1/sub`, you can do it like this:\n```python\nfrom fastapi import FastAPI\n\nsubapi = FastAPI()\n\n# ... define routes for subapi ...\n\napp = FastAPI()\n\napp.mount(\"/api/v1/sub\", subapi)\n```\nFastAPI will internally use the `root_path` smartly, so it will just work [1].\n\nHowever, when using `APIRouter`, it's a bit more complex. `APIRouter` has a `mount` method, but it's not optimized to work with sub-applications. The best practice is to mount your sub-application directly to the main FastAPI instance using the full path prefix [3].\n\nIt's also worth noting that when you mount a sub-application, the sub-application will have its `root_path` correctly set [1]. \n\nYou can also use `root_path` to configure your application when you have a proxy with a stripped path prefix [2]. \n\nFor example, if you have a proxy that adds a path prefix to your application, you can use `root_path` to handle this case. \n\nBut if you are trying to mount a sub-application under an `APIRouter`, it is a known limitation and you might need to reconsider your design [3]. \n\nYou can find more information in the documentation about [mounting sub-applications](https://fastapi.tiangolo.com/advanced/sub-applications/) and [using a proxy with a stripped path prefix](https://fastapi.tiangolo.com/advanced/behind-a-proxy/).",
|
| 165 |
+
"c1bd7ef1f9271cdf": "```json\n{\n \"faithful\": false,\n \"correct\": false\n}\n```",
|
| 166 |
+
"d0d0425aad0bc303": "FastAPI runs regular `def` endpoints without blocking the event loop by running them in a threadpool using `await run_in_threadpool(dependant.call, **values)` in the `run_endpoint_function` [4].\n\nThis allows the event loop to continue running other tasks while the endpoint function is executing, preventing blocking and ensuring that the application remains responsive.\n\nAdditionally, when needing to mix blocking and async code, it's recommended to use Asyncer, which provides a way to run blocking code inside async functions and vice versa using `asyncify()` and `syncify()` [5].",
|
| 167 |
+
"9ba358b72728a34d": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 168 |
+
"e54c5e2d5de4f8ce": "You can manage configuration and settings with environment variables using Pydantic Settings. This allows you to handle settings or configurations for your application with the power of Pydantic models. You can use a dependency to simplify testing, utilize `.env` files, and apply `@lru_cache` to avoid reading the dotenv file repeatedly while allowing overrides during testing [1].\n\nFirst, install the required package:\n```console\n$ pip install pydantic-settings\n```\nor \n```console\n$ pip install \"fastapi[all]\"\n```\nThen, create a `Settings` object by importing `BaseSettings` from Pydantic and creating a sub-class. Declare class attributes with type annotations and possible default values. Pydantic will read environment variables in a case-insensitive way, convert and validate the data [4].\n\nFor example, you can define a `config.py` file with:\n```python\nfrom pydantic import BaseSettings\n\nclass Settings(BaseSettings):\n admin_email: str\n app_name: str\n```\nAnd use it in your `main.py` file:\n```python\nfrom config import Settings\n\nsettings = Settings()\nprint(settings.admin_email)\nprint(settings.app_name)\n```\nYou can also use `.env` files to store environment variables. Create a `.env` file with:\n```bash\nADMIN_EMAIL=\"deadpool@example.com\"\nAPP_NAME=\"ChimichangApp\"\n```\nThen, update your `config.py` file to read from the `.env` file:\n```python\nfrom pydantic import BaseSettings\n\nclass Settings(BaseSettings):\n env_file = \".env\"\n admin_email: str\n app_name: str\n```\nTo avoid reading the `.env` file for each request, use the `@lru_cache` decorator:\n```python\nfrom functools import lru_cache\n\n@lru_cache()\ndef get_settings():\n return Settings()\n``` [2]\n\nYou can run the server with environment variables:\n```console\n$ ADMIN_EMAIL=\"deadpool@example.com\" APP_NAME=\"ChimichangApp\" fastapi run main.py\n``` [4]",
|
| 169 |
+
"a947beae5eed76b6": "```\n{\n \"faithful\": false,\n \"correct\": true\n}\n```",
|
| 170 |
+
"32500c98ea5e65c1": "The `APIRouter` class is defined in the `fastapi/routing.py` file [3][5]. You can also import it directly from `fastapi` [2]:\n```python\nfrom fastapi import APIRouter\n```",
|
| 171 |
+
"80c01ebf828e21ed": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 172 |
+
"042a557c2e4a4e71": "The `OAuth2PasswordBearer` class is implemented in the `fastapi/security/oauth2.py` file [1].",
|
| 173 |
+
"3c11ed85940cf1e7": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 174 |
+
"2325f5f6b0381148": "The `jsonable_encoder` function is implemented in the `fastapi/encoders.py` file [3].",
|
| 175 |
+
"8aa97952f2b8a102": "```json\n{\n \"faithful\": true,\n \"correct\": false\n}\n```",
|
| 176 |
+
"a336a2afff2a7f5a": "The `BackgroundTasks` class is defined in `fastapi/background.py`. It comes directly from [`starlette.background`](https://www.starlette.dev/background/), but is imported into FastAPI so you can import it from `fastapi` [2][1].",
|
| 177 |
+
"62e9b40118ab6542": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 178 |
+
"befaa7e5b9d09975": "FastAPI's `HTTPException` is defined in `fastapi/exceptions.py` [1].",
|
| 179 |
+
"30cc83d413377b42": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 180 |
+
"67b28e6202c7fc35": "The main FastAPI application class is defined in `fastapi/applications.py` as `FastAPI` [2].\n\nHere's an example of creating a FastAPI instance:\n```python\nfrom fastapi import FastAPI\n\napp = FastAPI()\n```",
|
| 181 |
+
"971c1829df5a8414": "```json\n{\n \"faithful\": false,\n \"correct\": false\n}\n```",
|
| 182 |
+
"a19e25729146e8f6": "These parameter functions, including `Query`, `Path`, and `Body`, are defined as functions that return special classes, and they can be imported directly from `fastapi` [3]. \n\nThey are actually functions that return objects of subclasses of a common `Param` class, which is itself a subclass of Pydantic's `FieldInfo` class [2].",
|
| 183 |
+
"81712b5cb22198eb": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 184 |
+
"dc681003c4ebc90e": "The `TestClient` that FastAPI provides for testing is located in `fastapi.testclient`. You can import it directly as:\n```python\nfrom fastapi.testclient import TestClient\n```\nIt is also available as `starlette.testclient.TestClient`, as FastAPI provides the same `starlette.testclient` as `fastapi.testclient` for convenience [3][4].",
|
| 185 |
+
"1ca643277d8b258f": "```json\n{\n \"faithful\": true,\n \"correct\": true\n}\n```",
|
| 186 |
+
"2e4133ebbed89f60": "To connect FastAPI to MongoDB, you would typically use a MongoDB client library for Python, such as `pymongo` or an ORM like `mongoengine` or `sqlmodel` doesn't support MongoDB but `mongoengine` does.\n\nFirst, you need to install the required library. For example, with `pymongo`:\n\n```bash\npip install pymongo\n```\n\nThen, you can use the library to connect to your MongoDB instance.\n\nHere is a basic example of how you might connect to MongoDB and use it with FastAPI:\n\n```python\nfrom fastapi import FastAPI\nfrom pymongo import MongoClient\n\n# Connect to MongoDB\nclient = MongoClient('mongodb://localhost:27017/')\ndb = client[\"your_database_name\"]\ncollection = db[\"your_collection_name\"]\n\napp = FastAPI()\n\n# Now you can use the collection to perform operations\n@app.get(\"/items/\")\ndef read_items():\n items = collection.find()\n return list(items)\n```\n\nHowever, for a real application, you'd likely want to handle connections and potential errors more robustly, possibly using dependency injection for database access.\n\nIn a containerized environment, you would also need to ensure MongoDB is running as a separate container and that your FastAPI application container can connect to it [1].\n\nFor more complex scenarios, consider using an async MongoDB client like `motor` which is compatible with async/await syntax used in FastAPI:\n\n```bash\npip install motor\n```\n\n```python\nfrom fastapi import FastAPI\nfrom motor.motor_asyncio import AsyncIOMotorClient\n\nclient = AsyncIOMotorClient('mongodb://localhost:27017/')\ndb = client[\"your_database_name\"]\ncollection = db[\"your_collection_name\"]\n\napp = FastAPI()\n\n@app.get(\"/items/\")\nasync def read_items():\n items = await collection.find().to_list(length=100)\n return items\n```",
|
| 187 |
+
"04b609bd6f97a903": "I couldn't find this in the indexed codebase.",
|
| 188 |
+
"e3a18d09f0cc9b25": "To integrate Celery task queues with FastAPI, you can use Celery as a tool for running background tasks that don't require sharing memory or variables with the same FastAPI app. This approach is suitable for heavy background computations where you don't need to access FastAPI app variables.\n\nHere's a basic example of how you might structure this:\n\nFirst, ensure you have Celery and a message broker like Redis or RabbitMQ installed. For this example, let's assume you're using Redis.\n\nThen, you can set up Celery in your FastAPI application like so:\n\n```python\nfrom fastapi import FastAPI\nfrom celery import Celery\n\napp = FastAPI()\n\n# Configure Celery\ncelery = Celery(\"tasks\", broker=\"redis://localhost:6379/0\")\n\n@celery.task\ndef my_background_task():\n # Your background task code here\n print(\"Running in the background\")\n\n@app.post(\"/start-task/\")\ndef start_task():\n my_background_task.apply_async()\n return {\"message\": \"Task started in the background\"}\n```\n\nIn this example, `my_background_task` is a Celery task that runs in the background. When you POST to `/start-task/`, it triggers `my_background_task` to run asynchronously.\n\nFor more complex use cases, including detailed configurations and monitoring, refer to the Celery documentation and [FastAPI's background tasks documentation][1].\n\n[1]: https://fastapi.tiangolo.com/tutorial/background-tasks/",
|
| 189 |
+
"d2137628f12fc5e5": "You can use FastAPI with the Django ORM by running your FastAPI application with a WSGI server that supports Django, or by using a third-party library that provides integration between FastAPI and Django.\n\nHowever, I couldn't find an official guide or a simple example in the provided context chunks. But I can tell you that it is possible to integrate FastAPI with Django.\n\nFastAPI is an ASGI framework, while Django is a WSGI framework. You can use a library like [Daphne](https://github.com/django/daphne) which is an ASGI server built for Django Channels, or [Hypercorn](https://hypercorn.readthedocs.io/) which is an ASGI server compatible with HTTP/2 and Trio.\n\nTo integrate FastAPI with the Django ORM, you would typically:\n- Run your Django application with the Django development server or a WSGI server like Gunicorn.\n- Run your FastAPI application with an ASGI server like Uvicorn.\n- Use a database like PostgreSQL that both Django and FastAPI can connect to.\n\nHere is an example of how you might structure your project:\n```python\n# myproject/settings.py (Django)\nDATABASES = {\n 'default': {\n 'ENGINE': 'django.db.backends.postgresql',\n 'NAME': 'mydatabase',\n 'USER': 'mydatabaseuser',\n 'PASSWORD': 'mypassword',\n 'HOST': '127.0.0.1',\n 'PORT': '5432',\n }\n}\n\n# myapp/main.py (FastAPI)\nfrom fastapi import FastAPI\nfrom django.core.management import call_command\nimport os\n\n# Configure Django\nos.environ.setdefault('DJANGO_SETTINGS_MODULE', 'myproject.settings')\ncall_command('makemigrations')\ncall_command('migrate')\n\nfrom myapp.models import MyModel\n\napp = FastAPI()\n\n@app.get(\"/items/\")\ndef read_items():\n items = MyModel.objects.all()\n return [{\"id\": item.id, \"name\": item.name} for item in items]\n```\nKeep in mind that this is a simplified example and does not cover error handling, database connections and other important aspects.\n\nFor more information, you can check [docs/en/docs/deployment/manually.md][1] and [docs/en/docs/advanced/wsgi.md][4].",
|
| 190 |
+
"7d309358c1c04b1d": "I couldn't find this in the indexed codebase.",
|
| 191 |
+
"9b649898daa89d96": "I couldn't find this in the indexed codebase.",
|
| 192 |
+
"5451166579a6804a": "I couldn't find this in the indexed codebase."
|
| 193 |
}
|
src/ask.py
CHANGED
|
@@ -1,30 +1,24 @@
|
|
| 1 |
-
"""Ask a question about the codebase — the full RAG pipeline
|
| 2 |
|
| 3 |
Flow:
|
| 4 |
-
1.
|
| 5 |
-
|
| 6 |
-
|
| 7 |
-
|
| 8 |
|
| 9 |
Usage:
|
| 10 |
export GROQ_API_KEY=gsk_... # free key from console.groq.com
|
| 11 |
python src/ask.py "How do I return a custom status code?"
|
| 12 |
-
python src/ask.py --
|
|
|
|
| 13 |
"""
|
| 14 |
|
| 15 |
import argparse
|
| 16 |
import os
|
| 17 |
import sys
|
| 18 |
-
from pathlib import Path
|
| 19 |
|
| 20 |
-
import
|
| 21 |
-
from groq import Groq
|
| 22 |
-
from sentence_transformers import SentenceTransformer
|
| 23 |
|
| 24 |
-
DATA_DIR = Path("data")
|
| 25 |
-
EMBED_MODEL = "BAAI/bge-small-en-v1.5"
|
| 26 |
-
# BGE models retrieve better when queries carry this instruction prefix
|
| 27 |
-
QUERY_PREFIX = "Represent this sentence for searching relevant passages: "
|
| 28 |
LLM_MODEL = os.environ.get("GROQ_MODEL", "openai/gpt-oss-120b")
|
| 29 |
TOP_K = 5
|
| 30 |
|
|
@@ -38,27 +32,11 @@ Answer ONLY from the numbered context chunks provided. Rules:
|
|
| 38 |
- Be concise."""
|
| 39 |
|
| 40 |
|
| 41 |
-
def retrieve(question: str, k: int) -> list[dict]:
|
| 42 |
-
model = SentenceTransformer(EMBED_MODEL)
|
| 43 |
-
query_emb = model.encode(QUERY_PREFIX + question,
|
| 44 |
-
normalize_embeddings=True)
|
| 45 |
-
client = chromadb.PersistentClient(path=str(DATA_DIR / "chroma"))
|
| 46 |
-
collection = client.get_collection("chunks")
|
| 47 |
-
res = collection.query(query_embeddings=[query_emb.tolist()], n_results=k)
|
| 48 |
-
return [{
|
| 49 |
-
"text": doc,
|
| 50 |
-
"meta": meta,
|
| 51 |
-
"distance": dist,
|
| 52 |
-
} for doc, meta, dist in zip(res["documents"][0],
|
| 53 |
-
res["metadatas"][0],
|
| 54 |
-
res["distances"][0])]
|
| 55 |
-
|
| 56 |
-
|
| 57 |
def build_prompt(question: str, hits: list[dict]) -> str:
|
| 58 |
parts = []
|
| 59 |
for i, h in enumerate(hits, 1):
|
| 60 |
parts.append(f"[{i}] ({h['meta']['source_type']}: "
|
| 61 |
-
|
| 62 |
context = "\n\n---\n\n".join(parts)
|
| 63 |
return f"Context chunks:\n\n{context}\n\nQuestion: {question}"
|
| 64 |
|
|
@@ -67,6 +45,7 @@ def answer(question: str, hits: list[dict]) -> str:
|
|
| 67 |
api_key = os.environ.get("GROQ_API_KEY")
|
| 68 |
if not api_key:
|
| 69 |
sys.exit("Set GROQ_API_KEY first (free key at console.groq.com).")
|
|
|
|
| 70 |
client = Groq(api_key=api_key)
|
| 71 |
response = client.chat.completions.create(
|
| 72 |
model=LLM_MODEL,
|
|
@@ -83,16 +62,18 @@ def main() -> None:
|
|
| 83 |
parser = argparse.ArgumentParser()
|
| 84 |
parser.add_argument("question")
|
| 85 |
parser.add_argument("--k", type=int, default=TOP_K)
|
|
|
|
|
|
|
| 86 |
parser.add_argument("--show-chunks", action="store_true",
|
| 87 |
help="print retrieved chunks (debugging/learning)")
|
| 88 |
args = parser.parse_args()
|
| 89 |
|
| 90 |
-
hits =
|
| 91 |
|
| 92 |
if args.show_chunks:
|
| 93 |
for i, h in enumerate(hits, 1):
|
| 94 |
-
print(f"\n=== [{i}]
|
| 95 |
-
f"{h['meta']['
|
| 96 |
print(h["text"][:500])
|
| 97 |
print("\n" + "=" * 60)
|
| 98 |
|
|
|
|
| 1 |
+
"""Ask a question about the codebase — the full RAG pipeline.
|
| 2 |
|
| 3 |
Flow:
|
| 4 |
+
1. Retrieve the most relevant chunks (dense, hybrid, or hybrid+rerank —
|
| 5 |
+
see retrieval.py for how each mode works)
|
| 6 |
+
2. Hand those chunks to an LLM and have it answer USING ONLY THEM
|
| 7 |
+
3. Print the answer plus links to the sources
|
| 8 |
|
| 9 |
Usage:
|
| 10 |
export GROQ_API_KEY=gsk_... # free key from console.groq.com
|
| 11 |
python src/ask.py "How do I return a custom status code?"
|
| 12 |
+
python src/ask.py --mode dense "..." # baseline retrieval
|
| 13 |
+
python src/ask.py --show-chunks "..." # also print retrieved chunks
|
| 14 |
"""
|
| 15 |
|
| 16 |
import argparse
|
| 17 |
import os
|
| 18 |
import sys
|
|
|
|
| 19 |
|
| 20 |
+
from retrieval import retrieve as retrieve_chunks
|
|
|
|
|
|
|
| 21 |
|
|
|
|
|
|
|
|
|
|
|
|
|
| 22 |
LLM_MODEL = os.environ.get("GROQ_MODEL", "openai/gpt-oss-120b")
|
| 23 |
TOP_K = 5
|
| 24 |
|
|
|
|
| 32 |
- Be concise."""
|
| 33 |
|
| 34 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 35 |
def build_prompt(question: str, hits: list[dict]) -> str:
|
| 36 |
parts = []
|
| 37 |
for i, h in enumerate(hits, 1):
|
| 38 |
parts.append(f"[{i}] ({h['meta']['source_type']}: "
|
| 39 |
+
f"{h['meta']['path']})\n{h['text'][:2500]}")
|
| 40 |
context = "\n\n---\n\n".join(parts)
|
| 41 |
return f"Context chunks:\n\n{context}\n\nQuestion: {question}"
|
| 42 |
|
|
|
|
| 45 |
api_key = os.environ.get("GROQ_API_KEY")
|
| 46 |
if not api_key:
|
| 47 |
sys.exit("Set GROQ_API_KEY first (free key at console.groq.com).")
|
| 48 |
+
from groq import Groq
|
| 49 |
client = Groq(api_key=api_key)
|
| 50 |
response = client.chat.completions.create(
|
| 51 |
model=LLM_MODEL,
|
|
|
|
| 62 |
parser = argparse.ArgumentParser()
|
| 63 |
parser.add_argument("question")
|
| 64 |
parser.add_argument("--k", type=int, default=TOP_K)
|
| 65 |
+
parser.add_argument("--mode", default="hybrid_rerank",
|
| 66 |
+
choices=["dense", "hybrid", "hybrid_rerank", "dense_rerank"])
|
| 67 |
parser.add_argument("--show-chunks", action="store_true",
|
| 68 |
help="print retrieved chunks (debugging/learning)")
|
| 69 |
args = parser.parse_args()
|
| 70 |
|
| 71 |
+
hits = retrieve_chunks(args.question, k=args.k, mode=args.mode)
|
| 72 |
|
| 73 |
if args.show_chunks:
|
| 74 |
for i, h in enumerate(hits, 1):
|
| 75 |
+
print(f"\n=== [{i}] {h['meta']['path']} :: "
|
| 76 |
+
f"{h['meta']['symbol']} ===")
|
| 77 |
print(h["text"][:500])
|
| 78 |
print("\n" + "=" * 60)
|
| 79 |
|
src/eval.py
CHANGED
|
@@ -1,20 +1,19 @@
|
|
| 1 |
"""Evaluate the RAG pipeline against a hand-labeled question set.
|
| 2 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 3 |
Metrics:
|
| 4 |
recall@k - did any gold source appear in the top-k retrieved chunks?
|
| 5 |
MRR - 1/rank of the first gold hit (higher = ranked better)
|
| 6 |
refusal - (with --answers) did unanswerable questions get a refusal?
|
| 7 |
faithful/correct - (with --answers) LLM-as-judge on generated answers
|
| 8 |
|
| 9 |
-
|
| 10 |
-
|
| 11 |
-
|
| 12 |
-
LLM answers and judgments are cached in data/eval_cache.json so re-runs
|
| 13 |
-
are free and fast (important on Groq's free-tier rate limits).
|
| 14 |
-
|
| 15 |
-
Usage:
|
| 16 |
-
python src/eval.py # retrieval metrics only (no LLM, fast)
|
| 17 |
-
python src/eval.py --answers # + generation, refusal, judge metrics
|
| 18 |
"""
|
| 19 |
|
| 20 |
import argparse
|
|
@@ -23,17 +22,14 @@ import json
|
|
| 23 |
import time
|
| 24 |
from pathlib import Path
|
| 25 |
|
| 26 |
-
import
|
| 27 |
-
from sentence_transformers import SentenceTransformer
|
| 28 |
|
| 29 |
DATA_DIR = Path("data")
|
| 30 |
EVAL_SET = DATA_DIR / "eval_set.jsonl"
|
| 31 |
CACHE_FILE = DATA_DIR / "eval_cache.json"
|
| 32 |
-
EMBED_MODEL = "BAAI/bge-small-en-v1.5"
|
| 33 |
-
QUERY_PREFIX = "Represent this sentence for searching relevant passages: "
|
| 34 |
K = 5
|
| 35 |
REFUSAL_TEXT = "I couldn't find this in the indexed codebase"
|
| 36 |
-
SLEEP_BETWEEN_LLM_CALLS = 5
|
| 37 |
|
| 38 |
JUDGE_PROMPT = """\
|
| 39 |
You are grading a RAG system's answer. Given the question, the context the
|
|
@@ -66,13 +62,16 @@ def cache_key(*parts: str) -> str:
|
|
| 66 |
return hashlib.sha256("||".join(parts).encode()).hexdigest()[:16]
|
| 67 |
|
| 68 |
|
| 69 |
-
def is_gold_hit(
|
| 70 |
-
|
|
|
|
| 71 |
return any(g.lower() in haystack for g in gold)
|
| 72 |
|
| 73 |
|
| 74 |
def main() -> None:
|
| 75 |
parser = argparse.ArgumentParser()
|
|
|
|
|
|
|
| 76 |
parser.add_argument("--answers", action="store_true",
|
| 77 |
help="also generate answers and run the LLM judge")
|
| 78 |
parser.add_argument("--k", type=int, default=K)
|
|
@@ -80,36 +79,27 @@ def main() -> None:
|
|
| 80 |
|
| 81 |
items = [json.loads(line)
|
| 82 |
for line in EVAL_SET.read_text().splitlines() if line.strip()]
|
| 83 |
-
print(f"[eval] {len(items)} questions "
|
| 84 |
f"({sum(i['answerable'] for i in items)} answerable)")
|
| 85 |
|
| 86 |
-
model = SentenceTransformer(EMBED_MODEL)
|
| 87 |
-
collection = chromadb.PersistentClient(
|
| 88 |
-
path=str(DATA_DIR / "chroma")).get_collection("chunks")
|
| 89 |
cache = load_cache()
|
| 90 |
|
| 91 |
# -------- retrieval metrics --------
|
| 92 |
recalls, mrrs = [], []
|
| 93 |
-
retrieved_per_q = []
|
| 94 |
for item in items:
|
| 95 |
-
|
| 96 |
-
normalize_embeddings=True)
|
| 97 |
-
res = collection.query(query_embeddings=[emb.tolist()],
|
| 98 |
-
n_results=args.k)
|
| 99 |
-
hits = list(zip(res["ids"][0], res["metadatas"][0],
|
| 100 |
-
res["documents"][0]))
|
| 101 |
retrieved_per_q.append(hits)
|
| 102 |
-
|
| 103 |
if not item["answerable"]:
|
| 104 |
continue
|
| 105 |
-
rank = next((r for r,
|
| 106 |
-
if is_gold_hit(
|
| 107 |
recalls.append(1.0 if rank else 0.0)
|
| 108 |
mrrs.append(1.0 / rank if rank else 0.0)
|
| 109 |
if not rank:
|
| 110 |
print(f" [miss] {item['question']}")
|
| 111 |
|
| 112 |
-
print(f"\n=== Retrieval (k={args.k}) ===")
|
| 113 |
print(f"recall@{args.k}: {sum(recalls)/len(recalls):.2f} "
|
| 114 |
f"({int(sum(recalls))}/{len(recalls)})")
|
| 115 |
print(f"MRR: {sum(mrrs)/len(mrrs):.2f}")
|
|
@@ -119,40 +109,40 @@ def main() -> None:
|
|
| 119 |
return
|
| 120 |
|
| 121 |
# -------- generation + judge metrics --------
|
| 122 |
-
from ask import SYSTEM_PROMPT, build_prompt # reuse the real pipeline
|
| 123 |
import os
|
|
|
|
|
|
|
| 124 |
from groq import Groq
|
| 125 |
client = Groq(api_key=os.environ["GROQ_API_KEY"])
|
| 126 |
llm_model = os.environ.get("GROQ_MODEL", "openai/gpt-oss-120b")
|
|
|
|
| 127 |
|
| 128 |
def llm(prompt: str, system: str | None = None) -> str:
|
| 129 |
-
|
| 130 |
-
|
| 131 |
-
|
| 132 |
-
|
| 133 |
-
|
| 134 |
-
|
| 135 |
-
|
| 136 |
-
|
| 137 |
-
|
| 138 |
-
|
| 139 |
-
|
| 140 |
-
|
| 141 |
-
|
| 142 |
-
|
| 143 |
-
|
| 144 |
-
|
| 145 |
-
|
| 146 |
-
|
| 147 |
-
|
| 148 |
-
|
| 149 |
-
|
| 150 |
|
| 151 |
refusal_ok, faithful, correct = [], [], []
|
| 152 |
for item, hits in zip(items, retrieved_per_q):
|
| 153 |
-
|
| 154 |
-
for cid, meta, doc in hits]
|
| 155 |
-
ans = llm(build_prompt(item["question"], hit_dicts),
|
| 156 |
system=SYSTEM_PROMPT)
|
| 157 |
|
| 158 |
if not item["answerable"]:
|
|
@@ -162,7 +152,7 @@ def main() -> None:
|
|
| 162 |
print(f" [no refusal] {item['question']}")
|
| 163 |
continue
|
| 164 |
|
| 165 |
-
context = "\n\n".join(h["text"][:
|
| 166 |
verdict_raw = llm(JUDGE_PROMPT.format(
|
| 167 |
question=item["question"], context=context, answer=ans))
|
| 168 |
try:
|
|
@@ -177,7 +167,7 @@ def main() -> None:
|
|
| 177 |
if not verdict.get("correct"):
|
| 178 |
print(f" [incorrect] {item['question']}")
|
| 179 |
|
| 180 |
-
print("\n=== Generation ===")
|
| 181 |
if faithful:
|
| 182 |
print(f"faithful: {sum(faithful)/len(faithful):.2f}")
|
| 183 |
print(f"correct: {sum(correct)/len(correct):.2f}")
|
|
|
|
| 1 |
"""Evaluate the RAG pipeline against a hand-labeled question set.
|
| 2 |
|
| 3 |
+
Now supports retrieval modes for ablation runs:
|
| 4 |
+
python src/eval.py --mode dense
|
| 5 |
+
python src/eval.py --mode hybrid
|
| 6 |
+
python src/eval.py --mode hybrid_rerank
|
| 7 |
+
python src/eval.py --mode hybrid_rerank --answers
|
| 8 |
+
|
| 9 |
Metrics:
|
| 10 |
recall@k - did any gold source appear in the top-k retrieved chunks?
|
| 11 |
MRR - 1/rank of the first gold hit (higher = ranked better)
|
| 12 |
refusal - (with --answers) did unanswerable questions get a refusal?
|
| 13 |
faithful/correct - (with --answers) LLM-as-judge on generated answers
|
| 14 |
|
| 15 |
+
LLM answers/judgments are cached in data/eval_cache.json (keyed by model
|
| 16 |
+
and prompt), so re-runs only pay for what changed.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 17 |
"""
|
| 18 |
|
| 19 |
import argparse
|
|
|
|
| 22 |
import time
|
| 23 |
from pathlib import Path
|
| 24 |
|
| 25 |
+
from retrieval import retrieve
|
|
|
|
| 26 |
|
| 27 |
DATA_DIR = Path("data")
|
| 28 |
EVAL_SET = DATA_DIR / "eval_set.jsonl"
|
| 29 |
CACHE_FILE = DATA_DIR / "eval_cache.json"
|
|
|
|
|
|
|
| 30 |
K = 5
|
| 31 |
REFUSAL_TEXT = "I couldn't find this in the indexed codebase"
|
| 32 |
+
SLEEP_BETWEEN_LLM_CALLS = 5
|
| 33 |
|
| 34 |
JUDGE_PROMPT = """\
|
| 35 |
You are grading a RAG system's answer. Given the question, the context the
|
|
|
|
| 62 |
return hashlib.sha256("||".join(parts).encode()).hexdigest()[:16]
|
| 63 |
|
| 64 |
|
| 65 |
+
def is_gold_hit(hit: dict, gold: list[str]) -> bool:
|
| 66 |
+
meta = hit["meta"]
|
| 67 |
+
haystack = f"{hit['id']} {meta.get('path', '')} {meta.get('symbol', '')}".lower()
|
| 68 |
return any(g.lower() in haystack for g in gold)
|
| 69 |
|
| 70 |
|
| 71 |
def main() -> None:
|
| 72 |
parser = argparse.ArgumentParser()
|
| 73 |
+
parser.add_argument("--mode", default="dense",
|
| 74 |
+
choices=["dense", "hybrid", "hybrid_rerank", "dense_rerank"])
|
| 75 |
parser.add_argument("--answers", action="store_true",
|
| 76 |
help="also generate answers and run the LLM judge")
|
| 77 |
parser.add_argument("--k", type=int, default=K)
|
|
|
|
| 79 |
|
| 80 |
items = [json.loads(line)
|
| 81 |
for line in EVAL_SET.read_text().splitlines() if line.strip()]
|
| 82 |
+
print(f"[eval] mode={args.mode}, {len(items)} questions "
|
| 83 |
f"({sum(i['answerable'] for i in items)} answerable)")
|
| 84 |
|
|
|
|
|
|
|
|
|
|
| 85 |
cache = load_cache()
|
| 86 |
|
| 87 |
# -------- retrieval metrics --------
|
| 88 |
recalls, mrrs = [], []
|
| 89 |
+
retrieved_per_q = []
|
| 90 |
for item in items:
|
| 91 |
+
hits = retrieve(item["question"], k=args.k, mode=args.mode)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 92 |
retrieved_per_q.append(hits)
|
|
|
|
| 93 |
if not item["answerable"]:
|
| 94 |
continue
|
| 95 |
+
rank = next((r for r, h in enumerate(hits, 1)
|
| 96 |
+
if is_gold_hit(h, item["gold"])), None)
|
| 97 |
recalls.append(1.0 if rank else 0.0)
|
| 98 |
mrrs.append(1.0 / rank if rank else 0.0)
|
| 99 |
if not rank:
|
| 100 |
print(f" [miss] {item['question']}")
|
| 101 |
|
| 102 |
+
print(f"\n=== Retrieval (mode={args.mode}, k={args.k}) ===")
|
| 103 |
print(f"recall@{args.k}: {sum(recalls)/len(recalls):.2f} "
|
| 104 |
f"({int(sum(recalls))}/{len(recalls)})")
|
| 105 |
print(f"MRR: {sum(mrrs)/len(mrrs):.2f}")
|
|
|
|
| 109 |
return
|
| 110 |
|
| 111 |
# -------- generation + judge metrics --------
|
|
|
|
| 112 |
import os
|
| 113 |
+
|
| 114 |
+
from ask import SYSTEM_PROMPT, build_prompt
|
| 115 |
from groq import Groq
|
| 116 |
client = Groq(api_key=os.environ["GROQ_API_KEY"])
|
| 117 |
llm_model = os.environ.get("GROQ_MODEL", "openai/gpt-oss-120b")
|
| 118 |
+
print(f"[eval] llm={llm_model}")
|
| 119 |
|
| 120 |
def llm(prompt: str, system: str | None = None) -> str:
|
| 121 |
+
key = cache_key(llm_model, system or "", prompt)
|
| 122 |
+
if key in cache:
|
| 123 |
+
return cache[key]
|
| 124 |
+
messages = ([{"role": "system", "content": system}] if system else [])
|
| 125 |
+
messages.append({"role": "user", "content": prompt})
|
| 126 |
+
out = None
|
| 127 |
+
for attempt in range(4):
|
| 128 |
+
try:
|
| 129 |
+
out = client.chat.completions.create(
|
| 130 |
+
model=llm_model, messages=messages,
|
| 131 |
+
temperature=0.1).choices[0].message.content
|
| 132 |
+
break
|
| 133 |
+
except Exception as e:
|
| 134 |
+
print(f" [retry {attempt + 1}/4] {str(e)[:160]}")
|
| 135 |
+
time.sleep(30)
|
| 136 |
+
if out is None:
|
| 137 |
+
raise RuntimeError("LLM call failed 4 times; try again later")
|
| 138 |
+
cache[key] = out
|
| 139 |
+
save_cache(cache)
|
| 140 |
+
time.sleep(SLEEP_BETWEEN_LLM_CALLS)
|
| 141 |
+
return out
|
| 142 |
|
| 143 |
refusal_ok, faithful, correct = [], [], []
|
| 144 |
for item, hits in zip(items, retrieved_per_q):
|
| 145 |
+
ans = llm(build_prompt(item["question"], hits),
|
|
|
|
|
|
|
| 146 |
system=SYSTEM_PROMPT)
|
| 147 |
|
| 148 |
if not item["answerable"]:
|
|
|
|
| 152 |
print(f" [no refusal] {item['question']}")
|
| 153 |
continue
|
| 154 |
|
| 155 |
+
context = "\n\n".join(h["text"][:1200] for h in hits)
|
| 156 |
verdict_raw = llm(JUDGE_PROMPT.format(
|
| 157 |
question=item["question"], context=context, answer=ans))
|
| 158 |
try:
|
|
|
|
| 167 |
if not verdict.get("correct"):
|
| 168 |
print(f" [incorrect] {item['question']}")
|
| 169 |
|
| 170 |
+
print(f"\n=== Generation (mode={args.mode}) ===")
|
| 171 |
if faithful:
|
| 172 |
print(f"faithful: {sum(faithful)/len(faithful):.2f}")
|
| 173 |
print(f"correct: {sum(correct)/len(correct):.2f}")
|
src/retrieval.py
ADDED
|
@@ -0,0 +1,126 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Retrieval backends: dense, hybrid (dense + BM25), and hybrid + reranker.
|
| 2 |
+
|
| 3 |
+
Modes:
|
| 4 |
+
dense - embedding similarity only (the v0 baseline)
|
| 5 |
+
hybrid - dense + BM25 keyword search, fused with Reciprocal Rank
|
| 6 |
+
Fusion (RRF)
|
| 7 |
+
hybrid_rerank - hybrid to get ~20 candidates, then a cross-encoder
|
| 8 |
+
re-scores each (question, chunk) pair and keeps the best
|
| 9 |
+
|
| 10 |
+
Why BM25 helps this corpus: code questions contain exact identifiers
|
| 11 |
+
("APIRouter", "jsonable_encoder"). Embeddings blur those into meaning;
|
| 12 |
+
BM25 matches them literally. The two are complementary, and RRF merges
|
| 13 |
+
their rankings without needing to calibrate scores against each other.
|
| 14 |
+
|
| 15 |
+
Why the reranker helps: the embedding compares question and chunk as two
|
| 16 |
+
separate vectors. A cross-encoder reads them TOGETHER, token by token, so
|
| 17 |
+
it is much better at judging true relevance - but too slow to run on all
|
| 18 |
+
1300 chunks, which is why it only re-scores the top candidates.
|
| 19 |
+
"""
|
| 20 |
+
|
| 21 |
+
import json
|
| 22 |
+
import re
|
| 23 |
+
from pathlib import Path
|
| 24 |
+
|
| 25 |
+
import chromadb
|
| 26 |
+
from rank_bm25 import BM25Okapi
|
| 27 |
+
from sentence_transformers import CrossEncoder, SentenceTransformer
|
| 28 |
+
|
| 29 |
+
DATA_DIR = Path("data")
|
| 30 |
+
EMBED_MODEL = "BAAI/bge-small-en-v1.5"
|
| 31 |
+
RERANK_MODEL = "BAAI/bge-reranker-v2-m3"
|
| 32 |
+
QUERY_PREFIX = "Represent this sentence for searching relevant passages: "
|
| 33 |
+
CANDIDATES = 20 # how many candidates hybrid gathers before final cut
|
| 34 |
+
RRF_K = 60 # standard RRF constant
|
| 35 |
+
|
| 36 |
+
# Lazy singletons so models/indexes load once per process, not per query
|
| 37 |
+
_embedder = None
|
| 38 |
+
_reranker = None
|
| 39 |
+
_collection = None
|
| 40 |
+
_bm25 = None
|
| 41 |
+
_chunk_ids: list[str] = []
|
| 42 |
+
_chunk_by_id: dict[str, dict] = {}
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def _tokenize(text: str) -> list[str]:
|
| 46 |
+
return re.findall(r"[a-z0-9_]+", text.lower())
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def _load() -> None:
|
| 50 |
+
global _embedder, _collection, _bm25, _chunk_ids, _chunk_by_id
|
| 51 |
+
if _embedder is not None:
|
| 52 |
+
return
|
| 53 |
+
_embedder = SentenceTransformer(EMBED_MODEL)
|
| 54 |
+
_collection = chromadb.PersistentClient(
|
| 55 |
+
path=str(DATA_DIR / "chroma")).get_collection("chunks")
|
| 56 |
+
chunks = [json.loads(line) for line in
|
| 57 |
+
(DATA_DIR / "chunks.jsonl").read_text().splitlines()]
|
| 58 |
+
_chunk_ids = [c["id"] for c in chunks]
|
| 59 |
+
_chunk_by_id = {c["id"]: c for c in chunks}
|
| 60 |
+
_bm25 = BM25Okapi([_tokenize(c["text"]) for c in chunks])
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+
def _get_reranker() -> CrossEncoder:
|
| 64 |
+
global _reranker
|
| 65 |
+
if _reranker is None:
|
| 66 |
+
_reranker = CrossEncoder(RERANK_MODEL) # first run downloads ~1GB
|
| 67 |
+
return _reranker
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
def _dense_ids(question: str, n: int) -> list[str]:
|
| 71 |
+
emb = _embedder.encode(QUERY_PREFIX + question, normalize_embeddings=True)
|
| 72 |
+
res = _collection.query(query_embeddings=[emb.tolist()], n_results=n)
|
| 73 |
+
return res["ids"][0]
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
def _bm25_ids(question: str, n: int) -> list[str]:
|
| 77 |
+
scores = _bm25.get_scores(_tokenize(question))
|
| 78 |
+
ranked = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
|
| 79 |
+
return [_chunk_ids[i] for i in ranked[:n]]
|
| 80 |
+
|
| 81 |
+
|
| 82 |
+
def _rrf_fuse(rankings: list[list[str]]) -> list[str]:
|
| 83 |
+
"""Merge multiple ranked lists: score(id) = sum over lists of
|
| 84 |
+
1/(RRF_K + rank). Ids high in ANY list surface; high in BOTH win."""
|
| 85 |
+
scores: dict[str, float] = {}
|
| 86 |
+
for ranking in rankings:
|
| 87 |
+
for rank, cid in enumerate(ranking, 1):
|
| 88 |
+
scores[cid] = scores.get(cid, 0.0) + 1.0 / (RRF_K + rank)
|
| 89 |
+
return sorted(scores, key=scores.get, reverse=True)
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
def _to_hit(cid: str) -> dict:
|
| 93 |
+
c = _chunk_by_id[cid]
|
| 94 |
+
return {
|
| 95 |
+
"id": cid,
|
| 96 |
+
"text": c["text"],
|
| 97 |
+
"meta": {
|
| 98 |
+
"source_type": c["source_type"],
|
| 99 |
+
"path": c["path"],
|
| 100 |
+
"symbol": c["symbol"] or "",
|
| 101 |
+
"url": c["url"],
|
| 102 |
+
},
|
| 103 |
+
}
|
| 104 |
+
|
| 105 |
+
def _rerank(question: str, candidates: list[str]) -> list[str]:
|
| 106 |
+
pairs = [(question, _chunk_by_id[cid]["text"][:2000])
|
| 107 |
+
for cid in candidates]
|
| 108 |
+
scores = _get_reranker().predict(pairs)
|
| 109 |
+
return [cid for _, cid in sorted(zip(scores, candidates), reverse=True)]
|
| 110 |
+
|
| 111 |
+
def retrieve(question: str, k: int = 5, mode: str = "dense") -> list[dict]:
|
| 112 |
+
_load()
|
| 113 |
+
if mode == "dense":
|
| 114 |
+
ids = _dense_ids(question, k)
|
| 115 |
+
elif mode == "dense_rerank":
|
| 116 |
+
ids = _rerank(question, _dense_ids(question, CANDIDATES))[:k]
|
| 117 |
+
elif mode in ("hybrid", "hybrid_rerank"):
|
| 118 |
+
fused = _rrf_fuse([_dense_ids(question, CANDIDATES),
|
| 119 |
+
_bm25_ids(question, CANDIDATES)])
|
| 120 |
+
candidates = fused[:CANDIDATES]
|
| 121 |
+
if mode == "hybrid_rerank":
|
| 122 |
+
candidates = _rerank(question, candidates)
|
| 123 |
+
ids = candidates[:k]
|
| 124 |
+
else:
|
| 125 |
+
raise ValueError(f"unknown mode: {mode}")
|
| 126 |
+
return [_to_hit(cid) for cid in ids]
|