Ultralytics Code Assistant

A code assistant that answers questions about Ultralytics YOLO by retrieving the relevant source code and generating a grounded answer.

▶ Watch the 2-minute overview: what was built, how it was measured, and what failed along the way. (The video was generated with Claude, based on this repository.)

At a glance

The core, as the brief asks:

Brief	Implementation
Chat interface	Gradio app: answers stream in, with the code they used shown next to them
Code indexing	ultralytics/models, engine and data at release v8.4.175: 104 files, 1,532 chunks split along functions, methods and classes with Python's ast
Vector search	MongoDB Atlas Vector Search over open-source embeddings (BAAI/bge-small-en-v1.5, run locally)
Answer generation	A free OpenRouter LLM, asked to cite each statement as [n], a link to the exact lines on GitHub

Optional extras, each with a switch in the chat. The first two were measured on 40 real questions:

• Keyword search and a reranker (on by default): related code reaches the LLM for 37 of 40 questions instead of 32.
• A second embedding model, Qwen3-Embedding-8B (served by OpenRouter): 40 of 40 with the reranker off, 39 with it on. It needs some OpenRouter credit, about $0.0000002 per question.
• A paid LLM (openai/gpt-5-nano), with the free models as fallback.

Built without RAG frameworks: plain pymongo, httpx and sentence-transformers, managed with uv.

Contents: Quick start · Using it · Example questions · Design notes · Evaluation · Limitations · Future work · Development

Quick start

What you need

	Why
uv ≥ 0.13 (curl -LsSf https://astral.sh/uv/install.sh | sh) and git	uv installs Python 3.12 and every dependency at the versions pinned in uv.lock. You don't need Python installed.
A free MongoDB Atlas cluster (M0)	Stores the chunks and runs vector and keyword search. Use an empty cluster: M0 allows 3 search indexes, and this project uses all 3.
A free OpenRouter API key	Generates the answers and reranks the sources. A key with no credit allows 50 free-model requests per day, about 20 questions.

No GPU is needed. The app runs on CPU and needs about 4 GB of free disk for the environment and the models.

Supported systems: Linux, and macOS 14 (Sonoma) or newer on Apple Silicon, both tested from a fresh clone following this README (the Mac was a Mac mini M1). Windows is covered by the locked dependencies but hasn't been tested. Intel Macs can't install it, because PyTorch no longer publishes wheels for them.

Creating the Atlas cluster and getting the connection string (about 5 minutes)

1. A project. Sign up, or log in and create a new project (Projects → New Project). If you already use Atlas, use a new project: each project can have one free cluster, and starting empty avoids the 3-index limit.
2. A free cluster in it. Click Create, choose Free (M0), any provider and region, and Create Deployment. Leave "Preload sample dataset" unchecked. The cluster is ready in 1–3 minutes.
3. A database user and your IP address. Atlas asks for both right after creating the cluster. Keep the proposed username and password, and copy the password now, because it is shown only once. Then add your current IP address.
4. The connection string. Open Connect → Drivers → Python and copy the string; ignore the pip install hint. It looks like mongodb+srv://<db_username>:<db_password>@cluster0.xxxxx.mongodb.net/?.... Replace <db_username> and <db_password>, including the < >, with the database user from step 3. That is not your Atlas login. The result goes in .env as MONGODB_URI=..., with no < or > left.

If you lost the password, or Atlas didn't create a user: Security → Database Access → Add New Database User, with password authentication, an autogenerated password and the role "Read and write to any database". Avoid passwords with special characters (@ : / ? # %): they break the connection string. The autogenerated ones are safe.

You don't create any index by hand: the index command creates them.

Run it from scratch

# 1. Get the code
git clone https://github.com/DanielPoncelas/ultralytics-code-assistant.git
cd ultralytics-code-assistant

# 2. Add your credentials: MONGODB_URI (Atlas step 4 above) and OPENROUTER_API_KEY
cp .env.example .env        # then edit .env

# 3. Install Python 3.12 and the dependencies, PyTorch included (no GPU needed)
uv sync

# 4. Check the configuration and both services
uv run code-assistant check

# 5. Build the index (one-off; from about a minute to a few)
uv run code-assistant index

# 6. Start the chat, then open http://127.0.0.1:7860
#    (the first start downloads a ~130 MB model, so it takes longer)
uv run code-assistant app

Step 4 should print three ticks. If a line shows ✘, its message says which setting or service to fix:

✔ Settings loaded
✔ MongoDB Atlas reachable (server 8.0.x)
✔ OpenRouter key valid (free-model requests today: 50/50)

Step 5 clones Ultralytics at v8.4.175 and verifies the commit. It chunks the code and uploads it to Atlas twice, once per embedding model (bge-small and Qwen3-8B, see Model choices), and creates the vector and keyword indexes. The chunk vectors are committed in embeddings/, so nothing has to be embedded.

It takes from about a minute to a few: most of it is waiting for Atlas to build the search indexes ("waiting until it is queryable …"). It's done when it prints one ✔ Indexed … line per embedding model. Re-running it is safe: chunks are updated in place, never duplicated.

Step 6 first prints "Loading the embedding model and connecting to Atlas …". The first time, it downloads bge-small (about 130 MB, pinned to an exact revision) from Hugging Face, so this takes longer; later starts take a few seconds. The chat is ready when the http://127.0.0.1:7860 link appears. Each question is then embedded locally (on the CPU, or on the GPU of an Apple Silicon Mac), which is why the default setup costs nothing.

Reproducing the embeddings. Delete embeddings/ and run step 5 again to recompute the vectors: about 7 minutes on a laptop CPU for bge-small, and less than a cent of OpenRouter credit for Qwen3-8B. The files are used only if their fingerprint (every chunk's text plus the model revision) matches, so stale vectors are never uploaded.

If something goes wrong

Symptom	Fix
check can't reach MongoDB Atlas, or the error says SSL handshake failed / TLSV1_ALERT_INTERNAL_ERROR	That error usually means your IP address isn't allowed, for example after changing networks. Check MONGODB_URI: are <db_username> and <db_password> replaced by the database user (not your Atlas login)? Is your IP address in Security → Network Access (Add IP Address → Add Current IP Address)? Does the password contain special characters?
index fails while creating an index	The cluster already has other search indexes, and M0 allows 3. Use an empty cluster in a new project, or set ALT_EMBEDDING_MODEL= (empty) to skip the second embedding model, which then needs only 2 indexes.
"No model could answer right now"	The free models are busy, or the key has used its 50 daily requests. Wait a minute, or turn on the paid model in the chat (needs OpenRouter credit).
"qwen/qwen3-embedding-8b needs OpenRouter credits"	Pick BAAI/bge-small-en-v1.5 as the embedding model, or add credit to the key.

Using it

Chat (uv run code-assistant app). Type a question or click one of the examples.

• The answer streams in with [n] citations.
• The panel on the right shows the code it used, each piece linked to its lines on GitHub.
• Follow-up questions work: they are rewritten into a standalone search, shown as "Searched for: …".

Above the sources, three options:

Option	Default	What it changes
Embedding model for the search	BAAI/bge-small-en-v1.5	qwen/qwen3-embedding-8b retrieves better but needs OpenRouter credit (about $0.0000002 per question)
Rerank the sources	on	Off: plain vector order, faster but less accurate
Use openai/gpt-5-nano (paid)	off	On: a stronger LLM answers (about $0.0004 per answer), with the free models as fallback

"Narrow the search" restricts the search to a folder (models, engine, data) or a YOLO task (detect, segment, …).

Command line:

uv run code-assistant ask "How do I resume an interrupted training run?"   # answer with cited sources
uv run code-assistant search "close_mosaic" --task detect -k 3            # retrieval only, no LLM

Configuration. Everything is read from .env or environment variables, and only the credentials are required:

Variable	Default	Purpose
MONGODB_URI, OPENROUTER_API_KEY	–	Credentials
EMBEDDING_MODEL	BAAI/bge-small-en-v1.5	Default embedding model
ALT_EMBEDDING_MODEL	qwen/qwen3-embedding-8b	Second model offered in the chat; empty disables it
LLM_MODELS	Cohere → dots-3 → Laguna (:free)	Free models, tried in order
PAID_LLM_MODEL	openai/gpt-5-nano	Model behind the paid switch; empty hides it
RERANK_MODEL	nvidia/llama-nemotron-rerank-vl-1b-v2:free	Reranker; empty disables it
RETRIEVAL_K	6	Code units given to the LLM

Example questions

Five questions taken from real Ultralytics GitHub issues. They are also the example buttons in the chat.

Each was answered end to end several times, with both embedding models and the default settings (these runs are not committed). 21 of the 22 answers contained every key fact. The answer shown for the augmentation question presents steps that are off by default as applied; see the note in docs/examples.md.

Question (source issue)	What the assistant answers
How do the lr0 and lrf hyperparameters work? Is the learning rate decayed linearly or with a cosine schedule? (#2849)	lr0 is the initial learning rate, and lr0 * lrf the final one. The decay is linear by default; a cosine (one-cycle) schedule is used only with cos_lr=True. It quotes BaseTrainer._setup_scheduler.
What does close_mosaic do, and at which point of training is mosaic augmentation switched off? (#14886)	Mosaic, copy-paste, mixup and cutmix are turned off for the last close_mosaic epochs. Once epoch >= epochs - close_mosaic, _close_dataloader_mosaic rebuilds the dataset's transforms.
Which data augmentations are applied by default when training a detection model? (#2412)	v8_transforms: Mosaic, RandomPerspective and CopyPaste, then MixUp, CutMix, Albumentations, HSV colour jitter and flips, each controlled by a hyperparameter. Limitation: the default values live in cfg/default.yaml, which is not indexed, so the answer can't tell that MixUp, CutMix, CopyPaste and vertical flips are off (0.0) by default.
How do I set my own class names for open-vocabulary detection with YOLO-World? (#20889)	model.set_classes([...]). It updates the class names (dropping the background entry) and resets the predictor, so the next prediction uses them.
How does the TensorRT export work? Does it go through ONNX first? (#1225)	Yes: export_engine exports ONNX first, then onnx2engine builds the .engine file, passing options such as INT8 quantization, dynamic shapes and the workspace size.

The full answers, unedited and with the sources they used, are in docs/examples.md. The questions the assistant gets wrong are described under Limitations.

Design notes

Architecture

Indexing (once: `index`)
  Ultralytics v8.4.175 (git clone, commit verified)
   └─► chunk with Python's ast: functions, methods, class and module summaries
        └─► embed: bge-small (local) and Qwen3-8B (OpenRouter); vectors shipped in embeddings/
             └─► MongoDB Atlas: one collection + vector index per embedding model, one keyword index

Answering (chat or `ask`)
  question ─► follow-up? rewrite it into a standalone search query
           ─► embed it with the selected model
           ─► Atlas $vectorSearch: top 20  +  Atlas Search (keywords): top 20
           ─► reranker (OpenRouter, free) keeps the best 6
           ─► prompt with 6 numbered sources ─► free LLM (Cohere → dots-3 → Laguna)
           ─► streamed answer with [n] citations linking to the exact lines on GitHub

How the code is chunked, and why

The code is split along its structure with Python's ast module, not into fixed-size windows:

Chunk	Content	Answers questions like
Function or method	The whole definition: decorators, signature, docstring, body	"What does BaseTrainer.save_model store?"
Class summary	Header, docstring, class attributes, and each method's signature and first docstring line	"What does DetectionTrainer override?"
Module summary	Docstring, imports, constants, and the signatures of top-level functions and classes	"What does ultralytics.data.augment contain?"

Why:

• Fixed-size windows cut functions in half and mix unrelated code.
• Whole files are too coarse: exporter.py has over 1,500 lines, too much for one vector.
• People ask about functions, methods and classes, and the summaries answer the "big picture" questions.
• The standard library's ast is enough for a pure-Python corpus.

Details that matter:

• Long units (over 4,000 characters, about 5% of them, e.g. the 444-line Exporter.__call__) are split at statement boundaries, never mid-statement, with a 3-line overlap.
• Comments, which ast ignores, are attached to the statement that follows them. A coverage check confirms that every line of code in the 104 files ends up in a chunk.
• The embedded text starts with a header naming the file, module, symbol, signature and part, because an isolated method body is ambiguous (self.model: which model?). The LLM gets the raw code.

Result: 1,532 chunks (1,072 methods, 214 class summaries, 140 functions and 106 module summaries; 79 long units split into parts). The median chunk is 927 characters.

Metadata extracted

Field	Used for
path, module, qualified_name, kind	Identifying and citing a chunk
start_line, end_line, github_url	Citations: a permalink to the exact lines at the pinned commit, so links never break
signature, docstring, decorators, parent_class, bases	Context for the LLM, and building the class summaries
folder, model_family (yolo, sam, rtdetr, …), task (detect, segment, …)	Pre-filters for the search, e.g. only task: segment
part, total_parts	Collapsing the parts of a long unit into one result
source (repo, tag, commit), embedding_model	Provenance; the app refuses to query an index built with another model
_id	A stable hash of path, symbol and part, so re-indexing updates instead of duplicating

Retrieval

1. Vector search: embed the question and run Atlas $vectorSearch with numCandidates = 20 × limit, MongoDB's guidance for near-exact recall. It matches an exact search on 99% of the top 5.
2. Keyword search: an Atlas Search index whose analyzer splits identifiers, so "nms" matches agnostic_nms and self.args.nms, and "trainer" matches DetectionTrainer. Together, the top 20 of each list contain the right code for 38 of 40 evaluation questions.
3. Reranking: a reranker reads each (question, code) pair and keeps the best 6, deduplicated so that a long function split into parts counts once. If it fails, for example when it is rate-limited, the search falls back to vector order instead of failing.

Model choices and rationale

Each choice was measured on the same 40 questions (see Evaluation), with the decision rule written down before the run.

Embeddings: BAAI/bge-small-en-v1.5 by default, qwen/qwen3-embedding-8b selectable.

• Four local open-source models were compared: Qwen3-Embedding-0.6B, Granite R2, GTE-ModernBERT and bge-small. The rule: best Hit@5 wins, unless a smaller model is within 3 points.
• bge-small won: one question behind Qwen-0.6B, ahead on Hit@1 and MRR, 18× smaller and 6.5× faster per query on CPU.
• Qwen3-Embedding-8B was added later. The brief asks for an open-source model, not a free one: its weights are Apache-2.0, and OpenRouter only hosts them. It reads 32k tokens, where bge-small truncates a third of the chunks at 512.
  • It gets related code into the top 6 for all 40 questions without the reranker (bge-small: 32), and for 39 with it (bge-small: 37).
  • It raises the share of key facts in the sources from 63% to 71–74%.
• bge-small stays the default because it is free and runs locally, so the app works with any OpenRouter key.

Reranker: nvidia/llama-nemotron-rerank-vl-1b-v2:free.

• Why hosted: a local cross-encoder took over a minute per question on CPU, so three OpenRouter rerankers were measured instead.
• The free one won: related code in the top 6 for 37 questions instead of 32, for about 0.6 s more per question.

LLM: cohere/north-mini-code:free, with dots-studio/dots-3-note-preview:free and poolside/laguna-s-2.1:free as fallbacks.

• How it was chosen: the free candidates answered all 46 questions from the same sources, and a judge from another model family graded them.
• Cohere won: it covered the most key facts (53.8%, 11 points ahead) and never declined a question it could answer.
• Why a chain: the fallbacks come from different providers, so one provider's outage doesn't stop the assistant.
• The paid switch (openai/gpt-5-nano) is there for when quality matters more than cost.

Answer generation

• The six sources are numbered in the prompt. The model must answer only from them and cite [n] per statement.
• When only part of the question is covered, it answers that part and says what is missing. It replies "I couldn't find this in the indexed code" only when nothing relevant was retrieved.
• The retrieved code is treated as untrusted data, so instructions hidden in it are ignored (covered by a prompt-injection test).
• If a model is rate-limited, times out or returns an empty answer, the next one in the chain answers. The status line says which model answered and which failed.

Trade-offs

Choice	Gained	Given up
bge-small as the default embedder	Free, local, 117 ms per query on CPU	It truncates long chunks; Qwen3-8B retrieves clearly better
Qwen3-8B behind a switch	The best retrieval measured	Needs OpenRouter credit and depends on one provider
Reranking on by default	5 more questions with related code in the top 6	About 0.6 s and one more free request per question
Free LLMs by default	Anyone can run it at no cost	They state about half the key facts, sometimes return empty answers, and get rate-limited
6 sources per question	The knee of the retrieval curve, about 6–7k prompt tokens	Answers that need code from many places
A pinned release	Reproducible results and permanent links	Doesn't follow new Ultralytics releases
Keyword search used only as reranker candidates	Catches exact identifiers	Mixed with vector results (RRF) without a reranker, it lost a question, so that mode stays off

Key decisions

Each decision is logged in DECISIONS.md, with the alternatives considered:

• D8: index a pinned release, verified by commit SHA.
• D9–D11: chunk along the syntax tree, split long units, add a contextual header.
• D12–D18: choose the embedding model with data; D28: offer Qwen3-8B.
• D6, D7, D19: an own OpenRouter client with ordered fallback, no SDK; D20: a cited, scoped, injection-aware prompt.
• D21, D23: grade answers with a judge validated by a human; choose the LLM.
• D24, D26: hybrid search tried and left off; reranker adopted.
• D14: ship the embeddings; D25: an optional paid model; D27: stricter retrieval metrics.

Evaluation

The full story (what was tried, in which order, with which rule, and how the judges were checked) is in EVALUATION.md. In short:

Question set. 46 questions from real GitHub issues: 40 that the indexed code can answer, and 6 that it can't. Before any run, each in-scope question was labelled with the code that answers it. Its 2–3 key facts and exact evidence lines were added after the embedding-model run, before any answer or fact was judged (eval/questions.toml).

Retrieval, from most lenient to strictest. Does the LLM get…

• Hit@6: at least one relevant piece of code?
• Facts@6: each key fact? Judged by an LLM that must quote the supporting line.
• All facts@6: every key fact, that is, the full answer?

Search	Hit@6	Facts@6	All facts@6
bge-small, vector only	80.0%	55.4%	35.0%
bge-small + reranker (default)	92.5%	62.9%	42.5%
Qwen3-8B, vector only	100%	71.4%	43.6%
Qwen3-8B + reranker	97.5%	73.9%	53.8%

Hit@6 alone would overstate quality. On "why do val() and predict() give different detections?", the search finds Model.val but not the two lines that answer the question (the conf defaults), so the answer misses the point.

Answers. The selected LLM states 53.8% of the key facts and has no false refusals. That was measured with vector-only retrieval, before the reranker; the example questions above were checked with the final setup.

Checking the judges. I reviewed samples blind, without seeing the judges' verdicts:

• Answer judge: it agreed with me on 25 of 25 key facts.
• Fact judge: 76% agreement, and every disagreement was the judge missing a fact that was there. So its numbers are a lower bound.

Reproduce:

uv run code-assistant eval retrieval       # embedding models  → eval/results/retrieval.md (~2 min on CPU)
uv run code-assistant eval search --facts  # search strategies → eval/results/search.md
uv run code-assistant eval answers         # LLMs + judge      → eval/results/answers.md

The last two use free OpenRouter requests (about 160, and 92 per LLM), so they need a key with at least $10 of credit purchased, which raises the free limit to 1,000 requests per day.

Limitations

• Answers spread across distant code. "Why do val() and predict() differ?" needs one line in Model.predict and one in BaseValidator.__init__. "How do I disable NMS?" needs a line inside setup_model. Neither embedding model retrieves those lines, so these answers are partial or wrong.
• The index covers only models/, engine/ and data/, as the brief asks. Defaults defined in cfg/default.yaml and helpers in utils/ aren't indexed, so questions about default values can be misread: once, the tuner's search bounds were presented as defaults.
• Free LLMs: they state about half the key facts, sometimes return empty answers or skip citations, and are rate-limited.
• A small evaluation: with 40 questions, only large differences are meaningful. The Qwen3-8B gain (8 questions fixed, none broken) is the clearest one.

Future work

What I would improve with more time

1. Answers that combine distant code. Let the LLM write extra search queries in the code's own vocabulary ("BaseValidator default conf"), or let it call the search as a tool several times before answering. This targets the failures above.
2. Finer retrieval inside long methods. Index small windows of long methods such as _do_train, and give the LLM the whole method when one of them matches.
3. Index cfg/default.yaml. It documents every argument and its default, which is where many real questions end up.
4. A larger, double-labelled evaluation set (100+ questions from issues and discussions) and a stronger judge, so that smaller differences become measurable.
5. Re-run the answer evaluation with the final retrieval and with paid models, to choose the default LLM on quality and cost.

How this would scale to production

• Indexing as a pipeline. Index every new Ultralytics release automatically (on a release webhook) into a versioned collection, and let users pick their version. Only re-embed the chunks that changed.
• Search. Move to a dedicated Atlas tier, with search nodes and room for more indexes, and quantize the vectors (int8) to cut memory. A self-hosted Qwen3-8B on a GPU removes the per-query dependency and cost.
• Serving. Put a stateless API (e.g. FastAPI) behind the UI, with autoscaling and concurrency. Use paid models with SLAs and provider routing, cache frequent questions, and add per-user rate limits and cost tracking.
• Quality and operations. Log each question with its retrieved code, model, latency and user feedback. Run the retrieval evaluation (about 2 minutes, no LLM) in CI on every change, and the answer evaluation nightly. Alert on fallback rates.
• Security. Use least-privilege database users (a separate reader and indexer) and a secrets manager, and keep treating retrieved text as untrusted.

Critical missing features

• Authentication and rate limiting: anyone who can reach the app spends the owner's API quota.
• Documentation and configuration files: answers come only from code, not from docs.ultralytics.com or the YAML configs.
• A single release: there is no way to ask about the version a user actually runs.
• No feedback loop: users can't flag a wrong answer, so quality can't be tracked in use.

Development

uv run pytest           # 149 tests, no network needed
uv run ruff check .     # lint
uv run ruff format .    # format

src/code_assistant/
  indexing/      fetch the pinned source, chunk with ast, embedding cache, upload pipeline
  evaluation/    question set, retrieval / search / answer evaluations, judges and human review
  embeddings.py  local (sentence-transformers) and OpenRouter-served embedding models
  store.py       MongoDB Atlas: vector and keyword search, indexes created from code
  retrieval.py   vector search, keyword candidates, reranking, fallbacks
  rerank.py      OpenRouter reranker client
  llm.py         OpenRouter chat client: fallback chain, streaming, retries
  prompt.py      prompt, citations, follow-up rewriting
  assistant.py   retrieve and generate
  app.py         Gradio chat
  cli.py         command line
eval/            labelled questions and committed results
embeddings/      precomputed chunk vectors per model
docs/            full example answers (examples.md) and the screenshot
EVALUATION.md    how quality was measured, what was tried, and the results
DECISIONS.md     decision log with alternatives and measurements