Merge pull request #43 from parea-ai/PAI-211-improve-trace-functional…

…ity-with-decorator

add traceable decorator and update logging

Makefile

@@ -21,7 +21,7 @@ poetry-remove:
 install:
 	poetry lock -n && poetry export --without-hashes > requirements.txt
 	poetry install -n
-	-poetry run mypy --install-types --non-interactive ./
+	poetry run mypy --install-types --non-interactive ./
 
 .PHONY: pre-commit-install
 pre-commit-install:

parea/api_client.py

@@ -1,11 +1,16 @@
 from typing import Any, Optional
 
+import os
+
 import httpx
+from dotenv import load_dotenv
+
+load_dotenv()
 
 
 class HTTPClient:
     _instance = None
-    base_url = "https://optimus-prompt-backend.vercel.app/api/parea/v1"
+    base_url = os.getenv("DEV_HOST") or "https://optimus-prompt-backend.vercel.app/api/parea/v1"
     api_key = None
 
     def __new__(cls, *args, **kwargs):

parea/client.py

@@ -3,12 +3,13 @@
 from attrs import asdict, define, field
 
 from parea.api_client import HTTPClient
-from parea.schemas.models import Completion, CompletionResponse, FeedbackRequest, LogRequest, UseDeployedPrompt, UseDeployedPromptResponse
+from parea.parea_logger import parea_logger
+from parea.schemas.models import Completion, CompletionResponse, FeedbackRequest, UseDeployedPrompt, UseDeployedPromptResponse
+from parea.utils.trace_utils import get_current_trace_id
 
 COMPLETION_ENDPOINT = "/completion"
 DEPLOYED_PROMPT_ENDPOINT = "/deployed-prompt"
 RECORD_FEEDBACK_ENDPOINT = "/feedback"
-LOG_ENDPOINT = "/log"
 
 
 @define
@@ -18,8 +19,10 @@ class Parea:
 
     def __attrs_post_init__(self):
         self._client.set_api_key(self.api_key)
+        parea_logger.set_client(self._client)
 
     def completion(self, data: Completion) -> CompletionResponse:
+        data.inference_id = get_current_trace_id()
         r = self._client.request(
             "POST",
             COMPLETION_ENDPOINT,
@@ -28,6 +31,7 @@ def completion(self, data: Completion) -> CompletionResponse:
         return CompletionResponse(**r.json())
 
     async def acompletion(self, data: Completion) -> CompletionResponse:
+        data.inference_id = get_current_trace_id()
         r = await self._client.request_async(
             "POST",
             COMPLETION_ENDPOINT,
@@ -65,20 +69,6 @@ async def arecord_feedback(self, data: FeedbackRequest) -> None:
             data=asdict(data),
         )
 
-    def log(self, data: LogRequest) -> None:
-        self._client.request(
-            "POST",
-            LOG_ENDPOINT,
-            data=asdict(data),
-        )
-
-    async def alog(self, data: LogRequest) -> None:
-        await self._client.request_async(
-            "POST",
-            LOG_ENDPOINT,
-            data=asdict(data),
-        )
-
 
 def gen_trace_id() -> str:
     """Generate a unique trace id for each chain of requests"""