add the code

context_manager.py

@@ -0,0 +1,133 @@
+from typing import Dict, Optional
+
+from schemas import CodeTaskType, CodeXRequest
+
+
+FOLLOW_UP_PATTERNS = [
+    "update this",
+    "modify this",
+    "change this",
+    "improve this",
+    "refactor this",
+    "fix this",
+    "explain this",
+    "review this",
+    "add validation",
+    "add error handling",
+    "add comments",
+    "optimize this",
+    "continue this",
+    "continue from above",
+    "same code",
+    "this code",
+    "that code",
+    "the above code",
+    "use the previous code",
+    "based on the previous code",
+]
+
+
+class SessionArtifact:
+    def __init__(
+        self,
+        code: Optional[str] = None,
+        task_type: Optional[CodeTaskType] = None,
+        language: Optional[str] = None,
+        framework: Optional[str] = None,
+        last_user_goal: Optional[str] = None,
+    ):
+        self.code = code
+        self.task_type = task_type
+        self.language = language
+        self.framework = framework
+        self.last_user_goal = last_user_goal
+
+
+class ContextManager:
+    def __init__(self):
+        self._store: Dict[str, SessionArtifact] = {}
+
+    def normalize_session_id(self, session_id: Optional[str]) -> str:
+        if session_id and str(session_id).strip():
+            return str(session_id).strip()
+        return "default"
+
+    def get_artifact(self, session_id: Optional[str]) -> Optional[SessionArtifact]:
+        normalized_session_id = self.normalize_session_id(session_id)
+        return self._store.get(normalized_session_id)
+
+    def save_artifact(
+        self,
+        session_id: Optional[str],
+        code: Optional[str],
+        task_type: Optional[CodeTaskType] = None,
+        language: Optional[str] = None,
+        framework: Optional[str] = None,
+        last_user_goal: Optional[str] = None,
+    ) -> None:
+        if not code or not str(code).strip():
+            return
+
+        normalized_session_id = self.normalize_session_id(session_id)
+        self._store[normalized_session_id] = SessionArtifact(
+            code=code.strip(),
+            task_type=task_type,
+            language=language.strip() if language else None,
+            framework=framework.strip() if framework else None,
+            last_user_goal=last_user_goal.strip() if last_user_goal else None,
+        )
+
+    def has_artifact(self, session_id: Optional[str]) -> bool:
+        artifact = self.get_artifact(session_id)
+        return bool(artifact and artifact.code and artifact.code.strip())
+
+    def is_follow_up_request(self, request: CodeXRequest) -> bool:
+        if request.code and request.code.strip():
+            return False
+
+        message = (request.message or "").strip().lower()
+        if not message:
+            return False
+
+        return any(pattern in message for pattern in FOLLOW_UP_PATTERNS)
+
+    def enrich_request_with_context(
+        self,
+        request: CodeXRequest,
+        session_id: Optional[str],
+    ) -> CodeXRequest:
+        if request.code and request.code.strip():
+            return request
+
+        artifact = self.get_artifact(session_id)
+        if not artifact or not artifact.code:
+            return request
+
+        if not self.is_follow_up_request(request):
+            return request
+
+        enriched_language = request.language or artifact.language
+        enriched_framework = request.framework or artifact.framework
+
+        previous_context_parts = []
+        if artifact.last_user_goal:
+            previous_context_parts.append(f"Previous User Goal:\n{artifact.last_user_goal}")
+        if request.previous_context and request.previous_context.strip():
+            previous_context_parts.append(request.previous_context.strip())
+
+        merged_previous_context = "\n\n".join(previous_context_parts) if previous_context_parts else request.previous_context
+
+        return CodeXRequest(
+            message=request.message,
+            mode=request.mode,
+            language=enriched_language,
+            code=artifact.code,
+            error_message=request.error_message,
+            framework=enriched_framework,
+            file_name=request.file_name,
+            previous_context=merged_previous_context,
+            use_retrieval=request.use_retrieval,
+        )
+
+
+context_manager = ContextManager()