From 45963e81f63f9413c4cea7dd47cb0242ddcd32e5 Mon Sep 17 00:00:00 2001 From: "sangsup.lee" Date: Thu, 18 Sep 2025 15:36:11 +0900 Subject: [PATCH] Add vllm feature for local network --- README.md | 10 ++++++++-- androidmeda.py | 17 ++++++++++++----- 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 55887d1..0e5b9ab 100644 --- a/README.md +++ b/README.md @@ -84,13 +84,19 @@ Follow steps here to download and run the model locally [github.com/ollama/ollam * **33B models:** Minimum 32GB RAM, 64GB recommended. Ensure your system (or WSL instance) has sufficient free memory. -To Run script with Ollama, +To Run script with Ollama, `python3 androidmeda.py --llm_provider ollama --llm_model llama3.2 -output_dir /tmp/ver/ -source_dir "input_dir1/ input_dir2/"` +**c. Using vLLM models** + +To Run script with vLLM, + +`python3 androidmeda.py --llm_provider vllm --llm_model your-model-name -output_dir /tmp/ver/ -source_dir "input_dir1/ input_dir2/"` + **Parameters -** -*-llm_provider* is the LLM provider of the model. e.g. google, anthropic, openaI, ollama +*-llm_provider* is the LLM provider of the model. e.g. google, anthropic, openai, ollama, vllm *-llm_model* is the LLM model to use, Gemini, Claude, ChatGPT are supported. You can get the model variants from here. [google](https://ai.google.dev/gemini-api/docs/models/gemini#model-variations) diff --git a/androidmeda.py b/androidmeda.py index f7d6afa..e51efe5 100644 --- a/androidmeda.py +++ b/androidmeda.py @@ -15,7 +15,7 @@ import anthropic _LLM_PROVIDER = flags.DEFINE_string( - 'llm_provider', None, 'LLM Provider to use e.g. google, openai, anthropic, ollama') + 'llm_provider', None, 'LLM Provider to use e.g. google, openai, anthropic, ollama, vllm') _LLM_MODEL = flags.DEFINE_string( 'llm_model', None, 'LLM Model to use e.g gemini-2.0-flash, gpt-4.1') _OUTPUT_DIR = flags.DEFINE_string( @@ -43,7 +43,7 @@ async def send_code_to_llm(system_instructions, files_data, llm_client=None): elif "ollama" in _LLM_PROVIDER.value: response = ollama.generate(model=_LLM_MODEL.value, format="json", prompt=complete_prompt) return response.response - elif "openai" in _LLM_PROVIDER.value: + elif "openai" in _LLM_PROVIDER.value or "vllm" in _LLM_PROVIDER.value: chat_completion = llm_client.chat.completions.create( messages=[ {"role": "system", "content": system_instructions}, @@ -176,15 +176,20 @@ async def main(argv: Sequence[str]) -> None: llm_client = None api_key = os.environ.get('API_KEY') - + api_base_url = os.environ.get('API_BASE_URL') + if _LLM_PROVIDER.value is None: raise app.UsageError( - f'Usage: Model provider is required e.g google, anthropic, openai, ollama' + f'Usage: Model provider is required e.g google, anthropic, openai, ollama, vllm' ) - elif "ollama" not in _LLM_PROVIDER.value and api_key is None: # Ollama does not require an API key + elif all(x not in _LLM_PROVIDER.value for x in ("ollama", "vllm")) and api_key is None: # Ollama and vLLM do not require an API key raise app.UsageError( f'Usage: {_LLM_PROVIDER.value} model requires an API key. Please set the API_KEY environment variable.' ) + elif "vllm" in _LLM_PROVIDER.value and api_base_url is None: + raise app.UsageError( + f'Usage: {_LLM_PROVIDER.value} model requires an API BASE URL. Please set the API_BASE_URL environment variable.' + ) elif _LLM_MODEL.value is None: raise app.UsageError( f'Usage: Model name is required e.g gemini-1.5-flash, gpt-4.1, llama3.2' @@ -199,6 +204,8 @@ async def main(argv: Sequence[str]) -> None: llm_client = anthropic.Anthropic(api_key=api_key) elif "ollama" in _LLM_PROVIDER.value: llm_client = None #We don't need to do anything + elif "vllm" in _LLM_PROVIDER.value: + llm_client = openai.OpenAI(api_key=api_key, base_url=api_base_url) else: raise ValueError(f"Unsupported LLM provider: {_LLM_PROVIDER.value}")