added milvus support and qol console logs for rag command #1720
Conversation
bmahabirbu
requested review from
rhatdan,
ericcurtin,
maxamillion,
swarajpande5,
jhjaggars,
cgruver and
engelmi
as code owners
Reviewer's GuideThis PR introduces Milvus support in the RAG workflow (documentation, installation, CLI, and config) and enhances console logging by adjusting stdout handling in the command runner. Entity relationship diagram for RAG format options including MilvuserDiagram
RAG_FORMATS {
string name
string description
}
RAG_FORMATS ||--o{ VECTOR_DB : supports
VECTOR_DB {
string type
}
VECTOR_DB {
qdrant
milvus
}
Class diagram for updated BaseConfig with Milvus supportclassDiagram
class BaseConfig {
int threads = -1
str port = str(DEFAULT_PORT)
str pull = "newer"
Literal["qdrant", "json", "markdown", "milvus"] rag_format = "qdrant"
SUPPORTED_RUNTIMES runtime = "llama.cpp"
str store
str temp = "0.8"
}
Flow diagram for command execution with enhanced stdout handlingflowchart TD
A[run_cmd called] --> B{stdout argument}
B -- Provided --> C[Use provided stdout]
B -- Not provided --> D[Default to subprocess.PIPE]
C --> E[Execute command]
D --> E[Execute command]
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey @bmahabirbu - I've reviewed your changes - here's some feedback:
- The change to run_cmd’s default stdout and the override in engine.run may have unintended side‐effects on other commands—consider isolating this to only the RAG workflow or restoring the previous defaults.
- You’ve added milvus as a CLI choice and in the Docker build, but there’s no actual milvus initialization or parameter handling—ensure the engine and CLI properly configure host/port, index settings, and connection logic for Milvus.
- The doc2rag and rag_framework scripts are listed in the diff but not updated—please clarify whether they need changes for Milvus support or remove those entries to avoid confusion.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The change to run_cmd’s default stdout and the override in engine.run may have unintended side‐effects on other commands—consider isolating this to only the RAG workflow or restoring the previous defaults.
- You’ve added milvus as a CLI choice and in the Docker build, but there’s no actual milvus initialization or parameter handling—ensure the engine and CLI properly configure host/port, index settings, and connection logic for Milvus.
- The doc2rag and rag_framework scripts are listed in the diff but not updated—please clarify whether they need changes for Milvus support or remove those entries to avoid confusion.
## Individual Comments
### Comment 1
<location> `ramalama/engine.py:165` </location>
<code_context>
dry_run(self.exec_args)
def run(self):
- run_cmd(self.exec_args)
+ run_cmd(self.exec_args, stdout=None)
def exec(self, stdout2null: bool = False, stderr2null: bool = False):
</code_context>
<issue_to_address>
Explicitly passing stdout=None changes subprocess output behavior.
Please verify that inheriting subprocess output is intentional and won't expose excessive or sensitive information in the console.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @bmahabirbu, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request expands the RAG command's capabilities by introducing Milvus as an alternative vector database backend. It also improves the user experience through enhanced console logging and refactors the internal architecture for better modularity. The rag_framework has been simplified by removing its HTTP API server component.
Highlights
- Milvus Vector Database Integration: I've added support for Milvus as a new vector database option for the RAG command. This includes installing
pymilvus, defining Milvus schemas, creating collections with indexes, and implementing data ingestion and querying logic. - Improved Console Logging: I've enhanced the user experience during document conversion by introducing real-time progress indicators. Error messages are now directed to
stderr. - Vector Database Abstraction and Refactoring: The interaction with vector databases (Qdrant and Milvus) has been refactored into dedicated classes within
rag_framework, allowing for dynamic selection of the vector store based on existing database files. - Removal of FastAPI Server from RAG Framework: The
rag_frameworkscript no longer includes the FastAPI server implementation, streamlining it to function as a CLI tool. - Hash Generation for Milvus Compatibility: The
generate_hashfunction indoc2raghas been updated to produceint64identifiers, which are required for Milvus primary keys.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
|
this will come in handy for the eventual Llamastack integration Although work will need to be done to use sentence transformer instead of fastembed |
There was a problem hiding this comment.
Code Review
This pull request introduces Milvus support and improves console logging for the rag command. I have identified some areas for improvement, including a critical issue with the serve command and some significant performance opportunities in the new Milvus implementation. Addressing these will enhance the robustness and efficiency of the new features.
| def serve_rag(vector_path): | ||
| rag = Rag(vector_path) | ||
| rag.serve() |
There was a problem hiding this comment.
The serve_rag function attempts to call rag.serve(), but the serve method was removed from the Rag class in this pull request (along with fastapi and uvicorn dependencies). This will cause a runtime AttributeError when a user runs the ramalama rag serve command.
To fix this, you should either:
- Re-implement the
servefunctionality if it's intended to be kept. - Remove the
servecommand entirely from the argument parser (inserve_parser) and delete theserve_ragfunction if this feature is no longer supported.
There was a problem hiding this comment.
fixed! Rag server wasnt used and will be implemented by proxy
| enable_dynamic_field=True, | ||
| ) | ||
| schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True) | ||
| schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=1000) |
There was a problem hiding this comment.
The max_length for the text field is hardcoded to 1000 characters. If the document chunker produces chunks larger than this, the text will be truncated upon insertion into Milvus, leading to silent data loss and potentially affecting the quality of RAG results.
Milvus supports a max_length of up to 65,535 for VARCHAR. I recommend increasing this limit significantly to prevent data truncation.
schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=65535)
There was a problem hiding this comment.
this is intensional to keep ram low but i will keep this in mind
| ) | ||
| # Chunk and add chunks to collection 1 by 1 | ||
| chunks, ids = self.chunk(docs) | ||
| for i, (chunk, id) in enumerate(zip(chunks, ids)): |
There was a problem hiding this comment.
Calling dmodel.embed() and smodel.embed() inside a loop for each chunk is highly inefficient, as fastembed is optimized for batch processing. This will be very slow for a large number of chunks.
A much more performant approach is to embed all chunks in a single batch before the loop. While this will use more memory to hold all embeddings at once, the speed improvement will be substantial. You can still insert them one by one or in smaller batches to keep memory usage low during the insertion phase.
# Batch-embed chunks for better performance
dense_embeddings = list(dmodel.embed(chunks))
sparse_embeddings_list = list(smodel.embed(chunks))
for i, (chunk, id) in enumerate(zip(chunks, ids)):
sparse_vector = sparse_embeddings_list[i].as_dict()
milvus_client.insert(
collection_name=collection_name,
data=[{
"id": id,
"text": chunk,
"sparse": sparse_vector,
"dense": dense_embeddings[i]
}]
)
# Flush every 100 records to reduce RAM usage
if (i + 1) % 100 == 0:
milvus_client.flush(collection_name=collection_name)
print(f"\rProcessed chunk {i+1}/{len(chunks)}", end='', flush=True)
There was a problem hiding this comment.
Very cool! Implemeneted and works!
| collection_name = COLLECTION_NAME | ||
| dmodel = TextEmbedding(model_name=EMBED_MODEL) | ||
| smodel = SparseTextEmbedding(model_name=SPARSE_MODEL) | ||
| test_embedding = next(dmodel.embed("This is a test")) |
There was a problem hiding this comment.
Instead of embedding a test string to find the embedding dimension, you can get it directly from the model object. Assuming a recent version of fastembed, you can use the .dim property, which is more efficient and readable.
embedding_dim = dmodel.dim
There was a problem hiding this comment.
'TextEmbedding' object has no attribute 'dim'
| metaprompt = f""" | ||
| You are an expert software architect. | ||
| Use the provided context and chat history to answer the question accurately and concisely. | ||
| If the answer is not explicitly stated, infer the most reasonable answer based on the available information. | ||
| If there is no relevant information, respond with "I don't know"—do not fabricate details. | ||
| You are an expert software architect. | ||
| Use the provided context and chat history to answer the question accurately and concisely. | ||
| If the answer is not explicitly stated, infer the most reasonable answer based on the available information. | ||
| If there is no relevant information, respond with "I don't know"—do not fabricate details. | ||
|
|
||
| ### Chat History: | ||
| {self.format_chat_history()} | ||
|
|
||
| ### Context: | ||
| {reranked_context.strip()} | ||
| ### Context: | ||
| {reranked_context.strip()} | ||
|
|
||
| ### Question: | ||
| {prompt.strip()} | ||
| ### Question: | ||
| {prompt.strip()} | ||
|
|
||
| ### Answer: | ||
| """ |
There was a problem hiding this comment.
The f-string for the metaprompt contains trailing whitespace on several lines. This can affect readability and is generally not a good practice. It's best to remove it for cleaner code.
metaprompt = f"""
You are an expert software architect.
Use the provided context and chat history to answer the question accurately and concisely.
If the answer is not explicitly stated, infer the most reasonable answer based on the available information.
If there is no relevant information, respond with \"I don't know\"—do not fabricate details.
### Chat History:
{self.format_chat_history()}
### Context:
{reranked_context.strip()}
### Question:
{prompt.strip()}
### Answer:
"""
42f4e6f to
cf88fe8
Compare
…ommand Signed-off-by: Brian <[email protected]>
| search_param_dense = { | ||
| "data": [dense_embedding], | ||
| "anns_field": "dense", | ||
| "param": {"metric_type": "IP", "params": {"nprobe": 10}}, |
There was a problem hiding this comment.
It would be better to define these magic numbers above.
| @@ -39,7 +39,8 @@ Convert documents into the following formats | |||
| | ------- | ---------------------------------------------------- | | |||
| | json | JavaScript Object Notation. lightweight format for exchanging data | | |||
| | markdown| Lightweight markup language using plain text editing | | |||
| | qdrant | Retrieval-Augmented Generation (RAG) Vector database | | |||
| | qdrant | Retrieval-Augmented Generation (RAG) Vector database Qdrant distribution | | |||
| | milvus | Retrieval-Augmented Generation (RAG) Vector database Milvus distribution | | |||
There was a problem hiding this comment.
These should be sorted alphabetically.
|
Couple of small comments that can be addressed in a separate PR, want to get this into the weekly release. |
Summary by Sourcery
Add support for Milvus as a new RAG data store option and improve console logging for the rag command
New Features:
Enhancements:
Build:
Documentation: