Merge branch 'master' of github.com:khoj-ai/lantern

app/settings.py

@@ -138,7 +138,7 @@
 DATABASES = {
     "default": {
         "ENGINE": "django.db.backends.postgresql",
-        "NAME": "postgres",
+        "NAME": os.getenv("DB_NAME", "postgres"),
         "USER": "postgres",
         "PASSWORD": os.getenv("POSTGRES_PASSWORD"),
         "HOST": os.environ.get("POSTGRES_HOST", "localhost"),

app/urls.py

@@ -24,6 +24,13 @@
     path("beta/", include("beta_product.urls")),
     path("auth/", include("user_manager.urls")),
     path("token/", views.obtain_auth_token, name="token"),
+    re_path(
+        r"^api/chat/.*", include("khoj_service.api_urls"), name="khoj_service_chat"
+    ),
+    # api/chat/ is a special case, because the response needs to be streamed just for this endpoint.
+    re_path(
+        r"^api/chat.*", include("khoj_service.chat_urls"), name="khoj_service_chat"
+    ),
     re_path(r"^api/.*", include("khoj_service.api_urls"), name="khoj_service_api"),
     path("", include("khoj_service.urls"), name="khoj_service"),
 ]

khoj_service/api_urls.py

@@ -1,6 +1,7 @@
 from django.urls import path
 from khoj_service import views
 
+
 urlpatterns = [
     path("", views.DefaultRedirectApiView.as_view()),
 ]

khoj_service/chat_urls.py

@@ -0,0 +1,7 @@
+from django.urls import path
+from khoj_service import views
+
+
+urlpatterns = [
+    path("", views.StreamingApiView.as_view()),
+]

khoj_service/constants.py

@@ -0,0 +1,3 @@
+from app.settings import DEBUG
+
+KHOJ_HOME = "http://localhost:3000" if DEBUG else "https://khoj.dev"

khoj_service/views.py

@@ -5,6 +5,10 @@
 from rest_framework.response import Response
 from rest_framework import status
 from khoj_service.models import RoutingTable
+from django.contrib.auth.models import AnonymousUser
+from django.shortcuts import redirect
+from django.http import StreamingHttpResponse
+from khoj_service.constants import KHOJ_HOME
 
 
 class DefaultRedirectApiView(APIView):
@@ -21,8 +25,7 @@ def get(self, request):
                 status=status.HTTP_404_NOT_FOUND,
             )
         khoj_response = requests.get(f"{service_url}{request_path}")
-        respond = HttpResponse(khoj_response)
-        return respond
+        return HttpResponse(khoj_response)
 
     def post(self, request):
         user = request.user
@@ -37,8 +40,41 @@ def post(self, request):
         khoj_response = requests.post(
             f"{service_url}{request_path}", json=body, headers=request.headers
         )
-        respond = HttpResponse(khoj_response)
-        return respond
+        return HttpResponse(khoj_response)
+
+
+class StreamingApiView(APIView):
+    permission_classes = [IsAuthenticated]
+    routing_table = RoutingTable.objects.all()
+
+    def _streaming_response(self, response):
+        aggregator = b""
+        for chunk in response.iter_content(100):
+            # This indicates that the compiled references are coming. We want to batch this response in a single chunk.
+            # Bear in mind it's possible that the compiled references indicator is split across two chunks. In this case, rendering will fail.
+            if (
+                len(aggregator) != 0
+                or str(chunk).find("### compiled references:") != -1
+            ):
+                aggregator += chunk
+                continue
+            yield chunk
+        yield aggregator
+
+    def get(self, request):
+        user = request.user
+        request_path = request.get_full_path()
+        service_url = self.routing_table.get_service_url_by_user(user)
+        if service_url is None:
+            return Response(
+                {"error": "user does not have a service with Khoj"},
+                status=status.HTTP_404_NOT_FOUND,
+            )
+        response = requests.get(f"{service_url}{request_path}", stream=True)
+        return StreamingHttpResponse(
+            streaming_content=self._streaming_response(response),
+            content_type=response.headers["Content-Type"],
+        )
 
 
 class RedirectToKhojAncilliaryAssets(APIView):
@@ -60,11 +96,13 @@ def get(self, request):
 
 
 class RedirectToKhojHomePage(APIView):
-    permission_classes = [IsAuthenticated]
     routing_table = RoutingTable.objects.all()
 
     def get(self, request):
         user = request.user
+        if isinstance(user, AnonymousUser):
+            return redirect(f"{KHOJ_HOME}/login/")
+
         service_url = self.routing_table.get_service_url_by_user(user)
         if service_url is None:
             # TODO: We probably want to redirect to a Sign-Up page here.