Skip to content
This repository was archived by the owner on Jun 5, 2025. It is now read-only.

Make weaviate conform with codegate logging #174

Merged
merged 1 commit into from
Dec 3, 2024
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
33 changes: 29 additions & 4 deletions src/codegate/storage/storage_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from weaviate.classes.query import MetadataQuery
from weaviate.embedded import EmbeddedOptions

from codegate.config import Config
from codegate.inference.inference_engine import LlamaCppInferenceEngine

logger = structlog.get_logger("codegate")
Expand All @@ -24,8 +25,32 @@
class StorageEngine:
def get_client(self, data_path):
try:
# Get current config
config = Config.get_config()

# Configure Weaviate logging
additional_env_vars = {
# Basic logging configuration
"LOG_FORMAT": config.log_format.value.lower(),
"LOG_LEVEL": config.log_level.value.lower(),
# Disable colored output
"LOG_FORCE_COLOR": "false",
# Configure JSON format
"LOG_JSON_FIELDS": "timestamp, level,message",
# Configure text format
"LOG_METHOD": config.log_format.value.lower(),
"LOG_LEVEL_IN_UPPER": "false", # Keep level lowercase like codegate format
# Disable additional fields
"LOG_GIT_HASH": "false",
"LOG_VERSION": "false",
"LOG_BUILD_INFO": "false",
}

client = weaviate.WeaviateClient(
embedded_options=EmbeddedOptions(persistence_data_path=data_path),
embedded_options=EmbeddedOptions(
persistence_data_path=data_path,
additional_env_vars=additional_env_vars,
),
)
return client
except Exception as e:
Expand All @@ -50,7 +75,7 @@ def __init__(self, data_path="./weaviate_data"):
try:
weaviate_client.close()
except Exception as e:
logger.info(f"Failed to close client: {str(e)}")
logger.error(f"Failed to close client: {str(e)}")
else:
logger.error("Could not find client, skipping schema setup.")

Expand All @@ -60,7 +85,7 @@ def setup_schema(self, client):
client.collections.create(
class_config["name"], properties=class_config["properties"]
)
logger.info(f"Weaviate schema for class {class_config['name']} setup complete.")
logger.info(f"Weaviate schema for class {class_config['name']} setup complete.")

async def search(self, query: str, limit=5, distance=0.3) -> list[object]:
"""
Expand Down Expand Up @@ -104,4 +129,4 @@ async def search(self, query: str, limit=5, distance=0.3) -> list[object]:
try:
weaviate_client.close()
except Exception as e:
logger.info(f"Failed to close client: {str(e)}")
logger.error(f"Failed to close client: {str(e)}")
6 changes: 4 additions & 2 deletions tests/test_storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,8 +35,10 @@ def mock_inference_engine():
@pytest.mark.asyncio
async def test_search(mock_weaviate_client, mock_inference_engine):
# Patch the LlamaCppInferenceEngine.embed method (not the entire class)
with patch("codegate.inference.inference_engine.LlamaCppInferenceEngine.embed",
mock_inference_engine.embed):
with patch(
"codegate.inference.inference_engine.LlamaCppInferenceEngine.embed",
mock_inference_engine.embed,
):

# Mock the WeaviateClient as before
with patch("weaviate.WeaviateClient", return_value=mock_weaviate_client):
Expand Down
Loading