Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 3 additions & 4 deletions doc/VectorCode-cli.txt
Original file line number Diff line number Diff line change
Expand Up @@ -771,7 +771,8 @@ will use that as the default project root for this process;
Note that:

1. For easier parsing, `--pipe` is assumed to be enabled in LSP mode;
2. At the time this only work with vectorcode setup that uses a **standalone ChromaDB server**, which is not difficult to setup using docker;
2. A `vectorcode.lock` file will be created in your `db_path` directory **if you’re using the bundled chromadb server**. Please do not delete it while a
vectorcode process is running;
3. The LSP server supports `vectorise`, `query` and `ls` subcommands. The other
subcommands may be added in the future.

Expand All @@ -789,9 +790,7 @@ features:
- `vectorise`vectorise files into a given project.

To try it out, install the `vectorcode[mcp]` dependency group and the MCP
server is available in the shell as `vectorcode-mcp-server`, and make sure
you’re using a |VectorCode-cli-standalone-chromadb-server| configured in the
|VectorCode-cli-json| via the `host` and `port` options.
server is available in the shell as `vectorcode-mcp-server`.

The MCP server entry point (`vectorcode-mcp-server`) provides some CLI options
that you can use to customise the default behaviour of the server. To view the
Expand Down
6 changes: 3 additions & 3 deletions doc/VectorCode.txt
Original file line number Diff line number Diff line change
Expand Up @@ -372,9 +372,9 @@ path to the executable) by calling `vim.lsp.config('vectorcode_server', opts)`.
minimal extra config required loading/unloading embedding models;
Progress reports.

Cons Heavy IO overhead because the Requires vectorcode-server; Only
embedding model and database works if you’re using a standalone
client need to be initialised ChromaDB server.
Cons Heavy IO overhead because the Requires vectorcode-server
embedding model and database
client need to be initialised
for every query.
-------------------------------------------------------------------------------
You may choose which backend to use by setting the |VectorCode-`setup`| option
Expand Down
9 changes: 4 additions & 5 deletions docs/cli.md
Original file line number Diff line number Diff line change
Expand Up @@ -696,8 +696,9 @@ will:
Note that:

1. For easier parsing, `--pipe` is assumed to be enabled in LSP mode;
2. At the time this only work with vectorcode setup that uses a **standalone
ChromaDB server**, which is not difficult to setup using docker;
2. A `vectorcode.lock` file will be created in your `db_path` directory __if
you're using the bundled chromadb server__. Please do not delete it while a
vectorcode process is running;
3. The LSP server supports `vectorise`, `query` and `ls` subcommands. The other
subcommands may be added in the future.

Expand All @@ -714,9 +715,7 @@ features:
- `vectorise`: vectorise files into a given project.

To try it out, install the `vectorcode[mcp]` dependency group and the MCP server
is available in the shell as `vectorcode-mcp-server`, and make sure you're using
a [standalone chromadb server](#chromadb) configured in the [JSON](#configuring-vectorcode)
via the `host` and `port` options.
is available in the shell as `vectorcode-mcp-server`.

The MCP server entry point (`vectorcode-mcp-server`) provides some CLI options
that you can use to customise the default behaviour of the server. To view the
Expand Down
2 changes: 1 addition & 1 deletion docs/neovim.md
Original file line number Diff line number Diff line change
Expand Up @@ -332,7 +332,7 @@ interface:
| Features | `default` | `lsp` |
|----------|-----------------------------------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------|
| **Pros** | Fully backward compatible with minimal extra config required | Less IO overhead for loading/unloading embedding models; Progress reports. |
| **Cons** | Heavy IO overhead because the embedding model and database client need to be initialised for every query. | Requires `vectorcode-server`; Only works if you're using a standalone ChromaDB server. |
| **Cons** | Heavy IO overhead because the embedding model and database client need to be initialised for every query. | Requires `vectorcode-server` |

You may choose which backend to use by setting the [`setup`](#setupopts) option `async_backend`,
and acquire the corresponding backend by the following API:
Expand Down
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ dependencies = [
"charset-normalizer>=3.4.1",
"json5",
"posthog<6.0.0",
"filelock>=3.15.0",
]
requires-python = ">=3.11,<3.14"
readme = "README.md"
Expand Down
29 changes: 29 additions & 0 deletions src/vectorcode/cli_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@

import json5
import shtab
from filelock import AsyncFileLock

from vectorcode import __version__

Expand Down Expand Up @@ -610,3 +611,31 @@ def config_logging(
handlers=handlers,
level=level,
)


class LockManager:
"""
A class that manages file locks that protects the database files in daemon processes (LSP, MCP).
"""

__locks: dict[str, AsyncFileLock]
singleton: Optional["LockManager"] = None

def __new__(cls) -> "LockManager":
if cls.singleton is None:
cls.singleton = super().__new__(cls)
cls.singleton.__locks = {}
return cls.singleton

def get_lock(self, path: str | os.PathLike) -> AsyncFileLock:
path = str(expand_path(str(path), True))
if os.path.isdir(path):
lock_file = os.path.join(path, "vectorcode.lock")
logger.info(f"Creating {lock_file} for locking.")
if not os.path.isfile(lock_file):
with open(lock_file, mode="w") as fin:
fin.write("")
path = lock_file
if self.__locks.get(path) is None:
self.__locks[path] = AsyncFileLock(path) # pyright: ignore[reportArgumentType]
return self.__locks[path]
110 changes: 82 additions & 28 deletions src/vectorcode/common.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,14 @@
import asyncio
import contextlib
import hashlib
import logging
import os
import socket
import subprocess
import sys
from typing import Any, AsyncGenerator
from asyncio.subprocess import Process
from dataclasses import dataclass
from typing import Any, AsyncGenerator, Optional
from urllib.parse import urlparse

import chromadb
Expand All @@ -16,7 +19,7 @@
from chromadb.config import APIVersion, Settings
from chromadb.utils import embedding_functions

from vectorcode.cli_utils import Config, expand_path
from vectorcode.cli_utils import Config, LockManager, expand_path

logger = logging.getLogger(name=__name__)

Expand Down Expand Up @@ -112,32 +115,6 @@ async def start_server(configs: Config):
return process


__CLIENT_CACHE: dict[str, AsyncClientAPI] = {}


async def get_client(configs: Config) -> AsyncClientAPI:
client_entry = configs.db_url
if __CLIENT_CACHE.get(client_entry) is None:
settings: dict[str, Any] = {"anonymized_telemetry": False}
if isinstance(configs.db_settings, dict):
valid_settings = {
k: v for k, v in configs.db_settings.items() if k in Settings.__fields__
}
settings.update(valid_settings)
parsed_url = urlparse(configs.db_url)
settings["chroma_server_host"] = parsed_url.hostname or "127.0.0.1"
settings["chroma_server_http_port"] = parsed_url.port or 8000
settings["chroma_server_ssl_enabled"] = parsed_url.scheme == "https"
settings["chroma_server_api_default_path"] = parsed_url.path or APIVersion.V2
settings_obj = Settings(**settings)
__CLIENT_CACHE[client_entry] = await chromadb.AsyncHttpClient(
settings=settings_obj,
host=str(settings_obj.chroma_server_host),
port=int(settings_obj.chroma_server_http_port or 8000),
)
return __CLIENT_CACHE[client_entry]


def get_collection_name(full_path: str) -> str:
full_path = str(expand_path(full_path, absolute=True))
hasher = hashlib.sha256()
Expand Down Expand Up @@ -261,3 +238,80 @@ async def list_collection_files(collection: AsyncCollection) -> list[str]:
or []
)
)


@dataclass
class _ClientModel:
client: AsyncClientAPI
is_bundled: bool = False
process: Optional[Process] = None


class ClientManager:
singleton: Optional["ClientManager"] = None
__clients: dict[str, _ClientModel]

def __new__(cls) -> "ClientManager":
if cls.singleton is None:
cls.singleton = super().__new__(cls)
cls.singleton.__clients = {}
return cls.singleton

@contextlib.asynccontextmanager
async def get_client(self, configs: Config, need_lock: bool = True):
project_root = str(expand_path(str(configs.project_root), True))
is_bundled = False
if self.__clients.get(project_root) is None:
process = None
if not await try_server(configs.db_url):
logger.info(f"Starting a new server at {configs.db_url}")
process = await start_server(configs)
is_bundled = True

self.__clients[project_root] = _ClientModel(
client=await self._create_client(configs),
is_bundled=is_bundled,
process=process,
)
lock = None
if self.__clients[project_root].is_bundled and need_lock:
lock = LockManager().get_lock(str(configs.db_path))
logger.debug(f"Locking {configs.db_path}")
await lock.acquire()
yield self.__clients[project_root].client
if lock is not None:
logger.debug(f"Unlocking {configs.db_path}")
await lock.release()

def get_processes(self) -> list[Process]:
return [i.process for i in self.__clients.values() if i.process is not None]

async def kill_servers(self):
termination_tasks: list[asyncio.Task] = []
for p in self.get_processes():
logger.info(f"Killing bundled chroma server with PID: {p.pid}")
p.terminate()
termination_tasks.append(asyncio.create_task(p.wait()))
await asyncio.gather(*termination_tasks)

async def _create_client(self, configs: Config) -> AsyncClientAPI:
settings: dict[str, Any] = {"anonymized_telemetry": False}
if isinstance(configs.db_settings, dict):
valid_settings = {
k: v for k, v in configs.db_settings.items() if k in Settings.__fields__
}
settings.update(valid_settings)
parsed_url = urlparse(configs.db_url)
settings["chroma_server_host"] = parsed_url.hostname or "127.0.0.1"
settings["chroma_server_http_port"] = parsed_url.port or 8000
settings["chroma_server_ssl_enabled"] = parsed_url.scheme == "https"
settings["chroma_server_api_default_path"] = parsed_url.path or APIVersion.V2
settings_obj = Settings(**settings)
return await chromadb.AsyncHttpClient(
settings=settings_obj,
host=str(settings_obj.chroma_server_host),
port=int(settings_obj.chroma_server_http_port or 8000),
)

def clear(self):
self.__clients.clear()
Loading