From c547b427022fb54d51910e9fabd2a979fa1d4958 Mon Sep 17 00:00:00 2001 From: valentimarco Date: Tue, 5 Nov 2024 16:26:47 +0100 Subject: [PATCH 1/5] feature(telemetry): add telemetry handler --- .env.example | 6 ++- compose.yml | 4 +- core/cat/env.py | 1 + core/cat/looking_glass/cheshire_cat.py | 7 ++++ core/cat/looking_glass/telemetry.py | 54 ++++++++++++++++++++++++++ core/cat/main.py | 2 +- core/cat/routes/embedder.py | 3 ++ core/cat/routes/llm.py | 4 ++ core/pyproject.toml | 1 + 9 files changed, 77 insertions(+), 5 deletions(-) create mode 100644 core/cat/looking_glass/telemetry.py diff --git a/.env.example b/.env.example index e8e07dba3..7e337d61c 100644 --- a/.env.example +++ b/.env.example @@ -1,4 +1,3 @@ - # Decide host and port for your Cat. Default will be localhost:1865 # CCAT_CORE_HOST=localhost # CCAT_CORE_PORT=1865 @@ -37,4 +36,7 @@ # CCAT_METADATA_FILE="cat/data/metadata.json" # Set container timezone -# CCAT_TIMEZONE=Europe/Rome \ No newline at end of file +# CCAT_TIMEZONE=Europe/Rome + +# Telemetry +CCAT_TELEMETRY=true \ No newline at end of file diff --git a/compose.yml b/compose.yml index 815bf1084..853bdf995 100644 --- a/compose.yml +++ b/compose.yml @@ -6,8 +6,8 @@ services: context: ./core container_name: cheshire_cat_core # Uncomment the two lines below to use your .env (see .env.example) - #env_file: - # - .env + env_file: + - .env ports: - ${CCAT_CORE_PORT:-1865}:80 - 5678:5678 # only for development purposes (take away in production) diff --git a/core/cat/env.py b/core/cat/env.py index 9506e2612..3fb72380e 100644 --- a/core/cat/env.py +++ b/core/cat/env.py @@ -21,6 +21,7 @@ def get_supported_env_variables(): "CCAT_JWT_EXPIRE_MINUTES": str(60 * 24), # JWT expires after 1 day "CCAT_HTTPS_PROXY_MODE": False, "CCAT_CORS_FORWARDED_ALLOW_IPS": "*", + "CCAT_TELEMETRY": True, } diff --git a/core/cat/looking_glass/cheshire_cat.py b/core/cat/looking_glass/cheshire_cat.py index 06b336bbf..a16995d4d 100644 --- a/core/cat/looking_glass/cheshire_cat.py +++ b/core/cat/looking_glass/cheshire_cat.py @@ -28,6 +28,7 @@ from cat.rabbit_hole import RabbitHole from cat.utils import singleton from cat import utils +from cat.env import get_env class Procedure(Protocol): name: str @@ -93,6 +94,12 @@ def __init__(self): # allows plugins to do something after the cat bootstrap is complete self.mad_hatter.execute_hook("after_cat_bootstrap", cat=self) + + # Telemetry + if get_env("CCAT_TELEMETRY") == "true": + from cat.looking_glass.telemetry import TelemetryHandler + log.info("Load Telemetry") + TelemetryHandler() def load_natural_language(self): """Load Natural Language related objects. diff --git a/core/cat/looking_glass/telemetry.py b/core/cat/looking_glass/telemetry.py new file mode 100644 index 000000000..faa9a686c --- /dev/null +++ b/core/cat/looking_glass/telemetry.py @@ -0,0 +1,54 @@ +from pydantic import BaseModel,Field +from typing import Optional +import psutil +from cat.looking_glass.white_rabbit import WhiteRabbit +from cat.utils import singleton +import httpx +from cat.log import log + + +class SystemData(BaseModel): + ram_gb: float =Field(frozen=True, default=round(psutil.virtual_memory().total / (1024.0 ** 3), 2)) + cpu_count: int = Field(frozen=True, default=psutil.cpu_count()) + + + +class TelemetryData(BaseModel): + country: Optional[str] = None + llm_model: Optional[str] = None + embedder_model: Optional[str] = None + system: SystemData + +@singleton +class TelemetryHandler: + + def __init__(self): + system=SystemData() + self.data = TelemetryData(system=system) + + def set_country(self,country: str): + # should be ISO3166 + self.data.country = country + + def set_llm_model(self, llm_model: str): + self.data.llm_model = llm_model + + def set_embedder_model(self, embedder_model: str): + self.data.embedder_model = embedder_model + + # @property + # def data(self) -> TelemetryData: + # return self.data + +def send_telemetry(): + telemetry = TelemetryHandler() + try: + log.info(f"Sending this chunk of data:{telemetry.data}") + # res = httpx.post("http://telemetry.cheshirecat.ai", data=telemetry.data) + # res.raise_for_status() + except httpx.HTTPStatusError as e: + log.error(f"Error when sending telemetry {e.response.status_code}") + + + +WhiteRabbit().schedule_interval_job(send_telemetry,hours=6) \ No newline at end of file diff --git a/core/cat/main.py b/core/cat/main.py index 28e2fcec4..dc6591e5f 100644 --- a/core/cat/main.py +++ b/core/cat/main.py @@ -29,7 +29,7 @@ host="0.0.0.0", port=80, use_colors=True, - log_level=get_env("CCAT_LOG_LEVEL").lower(), + log_level="debug", **debug_config, **proxy_pass_config, ) diff --git a/core/cat/routes/embedder.py b/core/cat/routes/embedder.py index 251a09d05..e4fa5ef9d 100644 --- a/core/cat/routes/embedder.py +++ b/core/cat/routes/embedder.py @@ -2,6 +2,7 @@ from cat.auth.connection import HTTPAuth from cat.auth.permissions import AuthPermission, AuthResource +from cat.looking_glass.telemetry import TelemetryHandler from fastapi import Request, APIRouter, Body, HTTPException, Depends from cat.factory.embedder import get_allowed_embedder_models, get_embedders_schemas @@ -143,6 +144,8 @@ def upsert_embedder_setting( ccat = request.app.state.ccat # reload llm and embedder of the cat ccat.load_natural_language() + + TelemetryHandler().set_embedder_model(languageEmbedderName) # crete new collections (different embedder!) try: ccat.load_memory() diff --git a/core/cat/routes/llm.py b/core/cat/routes/llm.py index 42f5dc8c1..be0e0ef5c 100644 --- a/core/cat/routes/llm.py +++ b/core/cat/routes/llm.py @@ -8,6 +8,7 @@ from cat.db import crud, models from cat.log import log from cat import utils +from cat.looking_glass.telemetry import TelemetryHandler router = APIRouter() @@ -127,6 +128,9 @@ def upsert_llm_setting( ccat = request.app.state.ccat # reload llm and embedder of the cat ccat.load_natural_language() + + TelemetryHandler().set_llm_model(languageModelName) + # crete new collections # (in case embedder is not configured, it will be changed automatically and aligned to vendor) # TODO: should we take this feature away? diff --git a/core/pyproject.toml b/core/pyproject.toml index 0c07d754c..c14214465 100644 --- a/core/pyproject.toml +++ b/core/pyproject.toml @@ -50,6 +50,7 @@ dependencies = [ "APScheduler==3.10.4", "ruff==0.4.7", "aiofiles==24.1.0", + "psutil==6.1.0", ] [tool.coverage.run] From 6174e8ff6e1269d84aa4a2780d007ddb2afeb8dc Mon Sep 17 00:00:00 2001 From: valentimarco Date: Tue, 5 Nov 2024 17:24:55 +0100 Subject: [PATCH 2/5] refactor(telemetry): remove singleton --- core/cat/looking_glass/cheshire_cat.py | 23 ++++++++++++++------ core/cat/looking_glass/telemetry.py | 29 +++++++++----------------- core/cat/routes/embedder.py | 2 -- core/cat/routes/llm.py | 3 --- 4 files changed, 27 insertions(+), 30 deletions(-) diff --git a/core/cat/looking_glass/cheshire_cat.py b/core/cat/looking_glass/cheshire_cat.py index a16995d4d..45adeed38 100644 --- a/core/cat/looking_glass/cheshire_cat.py +++ b/core/cat/looking_glass/cheshire_cat.py @@ -3,6 +3,7 @@ from typing_extensions import Protocol +from cat.looking_glass.telemetry import TelemetryHandler from langchain.base_language import BaseLanguageModel from langchain_core.messages import SystemMessage from langchain_core.runnables import RunnableLambda @@ -68,6 +69,13 @@ def __init__(self): # Start scheduling system self.white_rabbit = WhiteRabbit() + + # Telemetry + if get_env("CCAT_TELEMETRY") == "true": + from cat.looking_glass.telemetry import TelemetryHandler + log.info("Load Telemetry") + self.telemetry = TelemetryHandler() + self.white_rabbit.schedule_interval_job(self.telemetry.send_telemetry,seconds=6) # instantiate MadHatter (loads all plugins' hooks and tools) self.mad_hatter = MadHatter() @@ -94,12 +102,6 @@ def __init__(self): # allows plugins to do something after the cat bootstrap is complete self.mad_hatter.execute_hook("after_cat_bootstrap", cat=self) - - # Telemetry - if get_env("CCAT_TELEMETRY") == "true": - from cat.looking_glass.telemetry import TelemetryHandler - log.info("Load Telemetry") - TelemetryHandler() def load_natural_language(self): """Load Natural Language related objects. @@ -150,6 +152,8 @@ def load_language_model(self) -> BaseLanguageModel: selected_llm_config = crud.get_setting_by_name(name=selected_llm_class) try: llm = FactoryClass.get_llm_from_config(selected_llm_config["value"]) + if self.telemetry: + self.telemetry.set_llm_model(selected_llm_class) except Exception: import traceback @@ -193,6 +197,8 @@ def load_language_embedder(self) -> embedders.EmbedderSettings: embedder = FactoryClass.get_embedder_from_config( selected_embedder_config["value"] ) + if self.telemetry: + self.telemetry.set_embedder_model(selected_embedder_class) except AttributeError: import traceback @@ -433,3 +439,8 @@ def llm(self, prompt, *args, **kwargs) -> str: ) return output + + @property + def telemetryHandler(self): + if self.telemtry is None: + None diff --git a/core/cat/looking_glass/telemetry.py b/core/cat/looking_glass/telemetry.py index faa9a686c..6f6df59b3 100644 --- a/core/cat/looking_glass/telemetry.py +++ b/core/cat/looking_glass/telemetry.py @@ -1,3 +1,4 @@ +from uuid import UUID, uuid4 from pydantic import BaseModel,Field from typing import Optional import psutil @@ -11,15 +12,14 @@ class SystemData(BaseModel): ram_gb: float =Field(frozen=True, default=round(psutil.virtual_memory().total / (1024.0 ** 3), 2)) cpu_count: int = Field(frozen=True, default=psutil.cpu_count()) - - class TelemetryData(BaseModel): + telemetry_id: UUID = uuid4() country: Optional[str] = None llm_model: Optional[str] = None embedder_model: Optional[str] = None system: SystemData -@singleton + class TelemetryHandler: def __init__(self): @@ -36,19 +36,10 @@ def set_llm_model(self, llm_model: str): def set_embedder_model(self, embedder_model: str): self.data.embedder_model = embedder_model - # @property - # def data(self) -> TelemetryData: - # return self.data - -def send_telemetry(): - telemetry = TelemetryHandler() - try: - log.info(f"Sending this chunk of data:{telemetry.data}") - # res = httpx.post("http://telemetry.cheshirecat.ai", data=telemetry.data) - # res.raise_for_status() - except httpx.HTTPStatusError as e: - log.error(f"Error when sending telemetry {e.response.status_code}") - - - -WhiteRabbit().schedule_interval_job(send_telemetry,hours=6) \ No newline at end of file + def send_telemetry(self): + try: + log.info(f"Sending this chunk of data:{self.data}") + # res = httpx.post("http://telemetry.cheshirecat.ai", data=self.data) + # res.raise_for_status() + except httpx.HTTPStatusError as e: + log.error(f"Error when sending telemetry {e.response.status_code}") diff --git a/core/cat/routes/embedder.py b/core/cat/routes/embedder.py index e4fa5ef9d..a61c68b41 100644 --- a/core/cat/routes/embedder.py +++ b/core/cat/routes/embedder.py @@ -144,8 +144,6 @@ def upsert_embedder_setting( ccat = request.app.state.ccat # reload llm and embedder of the cat ccat.load_natural_language() - - TelemetryHandler().set_embedder_model(languageEmbedderName) # crete new collections (different embedder!) try: ccat.load_memory() diff --git a/core/cat/routes/llm.py b/core/cat/routes/llm.py index be0e0ef5c..179e8042d 100644 --- a/core/cat/routes/llm.py +++ b/core/cat/routes/llm.py @@ -128,9 +128,6 @@ def upsert_llm_setting( ccat = request.app.state.ccat # reload llm and embedder of the cat ccat.load_natural_language() - - TelemetryHandler().set_llm_model(languageModelName) - # crete new collections # (in case embedder is not configured, it will be changed automatically and aligned to vendor) # TODO: should we take this feature away? From 7202418bfb09fc485e7aa66eab1f1a3555c0b04c Mon Sep 17 00:00:00 2001 From: valentimarco Date: Tue, 5 Nov 2024 17:45:12 +0100 Subject: [PATCH 3/5] refactor: OPS... --- core/cat/main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/cat/main.py b/core/cat/main.py index dc6591e5f..28e2fcec4 100644 --- a/core/cat/main.py +++ b/core/cat/main.py @@ -29,7 +29,7 @@ host="0.0.0.0", port=80, use_colors=True, - log_level="debug", + log_level=get_env("CCAT_LOG_LEVEL").lower(), **debug_config, **proxy_pass_config, ) From 683e941ee176952ef66f2b479b85638ff3f0b469 Mon Sep 17 00:00:00 2001 From: valentimarco Date: Thu, 7 Nov 2024 21:19:02 +0100 Subject: [PATCH 4/5] refactor(telemetry)!: enable check moved in TelemetryHandler - also created a utility function to retrive cat version --- core/cat/looking_glass/cheshire_cat.py | 16 +++------ core/cat/looking_glass/telemetry.py | 50 ++++++++++++++++++++------ core/cat/routes/base.py | 8 ++--- core/cat/utils.py | 6 ++++ 4 files changed, 53 insertions(+), 27 deletions(-) diff --git a/core/cat/looking_glass/cheshire_cat.py b/core/cat/looking_glass/cheshire_cat.py index 45adeed38..14c5bf27b 100644 --- a/core/cat/looking_glass/cheshire_cat.py +++ b/core/cat/looking_glass/cheshire_cat.py @@ -29,7 +29,7 @@ from cat.rabbit_hole import RabbitHole from cat.utils import singleton from cat import utils -from cat.env import get_env + class Procedure(Protocol): name: str @@ -71,12 +71,8 @@ def __init__(self): self.white_rabbit = WhiteRabbit() # Telemetry - if get_env("CCAT_TELEMETRY") == "true": - from cat.looking_glass.telemetry import TelemetryHandler - log.info("Load Telemetry") - self.telemetry = TelemetryHandler() - self.white_rabbit.schedule_interval_job(self.telemetry.send_telemetry,seconds=6) - + self.telemetry = TelemetryHandler() + # instantiate MadHatter (loads all plugins' hooks and tools) self.mad_hatter = MadHatter() @@ -152,8 +148,7 @@ def load_language_model(self) -> BaseLanguageModel: selected_llm_config = crud.get_setting_by_name(name=selected_llm_class) try: llm = FactoryClass.get_llm_from_config(selected_llm_config["value"]) - if self.telemetry: - self.telemetry.set_llm_model(selected_llm_class) + self.telemetry.set_llm_model(selected_llm_class) except Exception: import traceback @@ -197,8 +192,7 @@ def load_language_embedder(self) -> embedders.EmbedderSettings: embedder = FactoryClass.get_embedder_from_config( selected_embedder_config["value"] ) - if self.telemetry: - self.telemetry.set_embedder_model(selected_embedder_class) + self.telemetry.set_embedder_model(selected_embedder_class) except AttributeError: import traceback diff --git a/core/cat/looking_glass/telemetry.py b/core/cat/looking_glass/telemetry.py index 6f6df59b3..5bb08b486 100644 --- a/core/cat/looking_glass/telemetry.py +++ b/core/cat/looking_glass/telemetry.py @@ -1,39 +1,67 @@ -from uuid import UUID, uuid4 -from pydantic import BaseModel,Field +from uuid import UUID +from pydantic import BaseModel, Field from typing import Optional import psutil from cat.looking_glass.white_rabbit import WhiteRabbit -from cat.utils import singleton import httpx from cat.log import log +from cat.env import get_env +from cat.utils import get_cat_version +from cat.db import crud, models class SystemData(BaseModel): - ram_gb: float =Field(frozen=True, default=round(psutil.virtual_memory().total / (1024.0 ** 3), 2)) + ram_gb: float = Field( + frozen=True, default=round(psutil.virtual_memory().total / (1024.0**3), 2) + ) cpu_count: int = Field(frozen=True, default=psutil.cpu_count()) + class TelemetryData(BaseModel): - telemetry_id: UUID = uuid4() + telemetry_id: UUID country: Optional[str] = None + version: str = get_cat_version() llm_model: Optional[str] = None embedder_model: Optional[str] = None system: SystemData class TelemetryHandler: - def __init__(self): - system=SystemData() - self.data = TelemetryData(system=system) - - def set_country(self,country: str): + self.enable: bool = get_env("CCAT_TELEMETRY") == "true" + if self.enable: + log.info("Load Telemetry") + + telemetry_settings = crud.get_setting_by_name("Telemetry") + if telemetry_settings: + # we use the setting_id as telemetry id + telemetry_id = telemetry_settings["setting_id"] + else: + setting = crud.create_setting( + models.Setting(name="Telemetry", category="telemetry", value={}) + ) + telemetry_id = setting["setting_id"] + + system = SystemData() + self.data = TelemetryData(telemetry_id=telemetry_id, system=system) + WhiteRabbit().schedule_interval_job(self.send_telemetry, seconds=6) + else: + log.info("Telemetry is disable") + + def set_country(self, country: str): + if not self.enable: + return # should be ISO3166 self.data.country = country def set_llm_model(self, llm_model: str): + if not self.enable: + return self.data.llm_model = llm_model - + def set_embedder_model(self, embedder_model: str): + if not self.enable: + return self.data.embedder_model = embedder_model def send_telemetry(self): diff --git a/core/cat/routes/base.py b/core/cat/routes/base.py index 2e6c431e5..d35d9a9b9 100644 --- a/core/cat/routes/base.py +++ b/core/cat/routes/base.py @@ -1,11 +1,12 @@ from fastapi import APIRouter, Depends, Body from fastapi.concurrency import run_in_threadpool from typing import Dict -import tomli + from cat.auth.permissions import AuthPermission, AuthResource from cat.auth.connection import HTTPAuth from cat.convo.messages import CatMessage +from cat.utils import get_cat_version router = APIRouter() @@ -16,10 +17,7 @@ async def status( stray=Depends(HTTPAuth(AuthResource.STATUS, AuthPermission.READ)), ) -> Dict: """Server status""" - with open("pyproject.toml", "rb") as f: - project_toml = tomli.load(f)["project"] - - return {"status": "We're all mad here, dear!", "version": project_toml["version"]} + return {"status": "We're all mad here, dear!", "version": get_cat_version()} @router.post("/message", response_model=CatMessage) diff --git a/core/cat/utils.py b/core/cat/utils.py index ba46745a0..7854eb935 100644 --- a/core/cat/utils.py +++ b/core/cat/utils.py @@ -12,6 +12,7 @@ from langchain_core.output_parsers import JsonOutputParser from langchain_core.prompts import PromptTemplate from langchain_core.utils import get_colored_text +import tomli from cat.log import log from cat.env import get_env @@ -243,6 +244,11 @@ def langchain_log_output(langchain_output, title): return langchain_output +def get_cat_version() -> str: + with open("pyproject.toml", "rb") as f: + project_toml = tomli.load(f)["project"] + return project_toml["version"] + # This is our masterwork during tea time class singleton: instances = {} From 2438ff475e6a033523d64b2fcb6a954a0d74320d Mon Sep 17 00:00:00 2001 From: valentimarco Date: Thu, 7 Nov 2024 21:21:07 +0100 Subject: [PATCH 5/5] lint: remove unused class --- core/cat/routes/embedder.py | 1 - core/cat/routes/llm.py | 1 - 2 files changed, 2 deletions(-) diff --git a/core/cat/routes/embedder.py b/core/cat/routes/embedder.py index a61c68b41..251a09d05 100644 --- a/core/cat/routes/embedder.py +++ b/core/cat/routes/embedder.py @@ -2,7 +2,6 @@ from cat.auth.connection import HTTPAuth from cat.auth.permissions import AuthPermission, AuthResource -from cat.looking_glass.telemetry import TelemetryHandler from fastapi import Request, APIRouter, Body, HTTPException, Depends from cat.factory.embedder import get_allowed_embedder_models, get_embedders_schemas diff --git a/core/cat/routes/llm.py b/core/cat/routes/llm.py index 179e8042d..42f5dc8c1 100644 --- a/core/cat/routes/llm.py +++ b/core/cat/routes/llm.py @@ -8,7 +8,6 @@ from cat.db import crud, models from cat.log import log from cat import utils -from cat.looking_glass.telemetry import TelemetryHandler router = APIRouter()