Dynavera/apps/mlstore/viewsets.py

from rest_framework.viewsets import ModelViewSet
from rest_framework.permissions import IsAuthenticated
from rest_framework.decorators import action
from rest_framework.response import Response
from rest_framework import status
from .models import Agent, AgentRun, AgentEvent
from .serializers import AgentSerializer, AgentRunSerializer, AgentEventSerializer
from . import services
from apps.orgs.models import Role


class AgentViewSet(ModelViewSet):
    queryset = Agent.objects.all()
    serializer_class = AgentSerializer
    permission_classes = [IsAuthenticated]
    lookup_field = 'uuid'


class AgentRunViewSet(ModelViewSet):
    queryset = AgentRun.objects.all()
    serializer_class = AgentRunSerializer
    permission_classes = [IsAuthenticated]
    lookup_field = 'uuid'

    def get_queryset(self):
        return AgentRun.objects.filter(user=self.request.user)

    @action(detail=True, methods=['get'], url_path='events')
    def events(self, request, uuid=None):
        run = self.get_object()
        events = AgentEvent.objects.filter(execution=run).order_by('timestamp')
        serializer = AgentEventSerializer(events, many=True)
        return Response(serializer.data)

    @action(detail=False, methods=['post'], url_path='retrieve-context')
    def retrieve_context(self, request):
        """Retrieve context documents from RAG using semantic similarity.
        
        Request body:
        {
            "query": "search query text",
            "role_uuid": "role-uuid",
            "top_k": 5,  # optional, default 5
            "similarity_threshold": 0.5  # optional, default 0.5
        }
        
        Returns:
        {
            "query": "search query text",
            "context": "formatted context string with sources",
            "documents": [
                {
                    "id": 123,
                    "content": "chunk text",
                    "similarity": 0.87,
                    "source": "filename.pdf",
                    "chunk_index": 0
                },
                ...
            ]
        }
        """
        query = request.data.get('query', '').strip()
        role_uuid = request.data.get('role_uuid', '').strip()
        top_k = request.data.get('top_k', 5)
        similarity_threshold = request.data.get('similarity_threshold', 0.5)
        
        if not query:
            return Response(
                {"error": "query is required"},
                status=status.HTTP_400_BAD_REQUEST
            )
        
        if not role_uuid:
            return Response(
                {"error": "role_uuid is required"},
                status=status.HTTP_400_BAD_REQUEST
            )
        
        try:
            # Validate role exists and user has access
            role = Role.objects.get(uuid=role_uuid)
            # You can add additional permission checks here if needed
            
            # Search for similar documents
            results = services.search_similar_documents(
                query=query,
                role_uuid=role_uuid,
                top_k=top_k,
                similarity_threshold=similarity_threshold
            )
            
            # Format response
            documents = []
            for doc, similarity in results:
                documents.append({
                    "id": doc.id,
                    "uuid": str(doc.uuid),
                    "content": doc.content,
                    "similarity": float(similarity),
                    "source": doc.training_file.file_name if doc.training_file else "unknown",
                    "chunk_index": doc.chunk_index,
                    "metadata": doc.metadata,
                })
            
            # Get formatted context string
            context = services.get_context_for_query(
                query=query,
                role_uuid=role_uuid,
                top_k=top_k,
                similarity_threshold=similarity_threshold
            )
            
            return Response({
                "query": query,
                "role_uuid": role_uuid,
                "num_results": len(documents),
                "context": context,
                "documents": documents,
            })
            
        except Role.DoesNotExist:
            return Response(
                {"error": f"Role with UUID {role_uuid} not found"},
                status=status.HTTP_404_NOT_FOUND
            )
        except ValueError as e:
            return Response(
                {"error": str(e)},
                status=status.HTTP_400_BAD_REQUEST
            )
        except Exception as e:
            import logging
            logging.exception("Error retrieving context")
            return Response(
                {"error": "Failed to retrieve context", "detail": str(e)},
                status=status.HTTP_500_INTERNAL_SERVER_ERROR
            )