1  import json
  2  import logging
  3  from uuid import uuid4
  4  
  5  from fastapi import HTTPException
  6  
  7  from restai.database import DBWrapper
  8  from restai.models.models import ChatModel, QuestionModel, User
  9  from restai.project import Project
 10  from restai.projects.base import ProjectBase
 11  from restai.projects.block_interpreter import BlockInterpreter
 12  
 13  logger = logging.getLogger(__name__)
 14  
 15  
 16  class Block(ProjectBase):
 17  
 18      def _get_workspace(self, project: Project) -> dict:
 19          workspace = None
 20          if project.props.options:
 21              opts = project.props.options
 22              if hasattr(opts, "blockly_workspace"):
 23                  workspace = opts.blockly_workspace
 24              elif isinstance(opts, dict):
 25                  workspace = opts.get("blockly_workspace")
 26          if not workspace:
 27              raise HTTPException(
 28                  status_code=400,
 29                  detail="No block workspace configured for this project.",
 30              )
 31          return workspace
 32  
 33      async def chat(self, project: Project, chat_model: ChatModel, user: User, db: DBWrapper):
 34          from restai.agent2.memory import get_session, save_session
 35          from restai.agent2.types import Message, TextBlock, user_text_message
 36  
 37          workspace = self._get_workspace(project)
 38          chat_id = chat_model.id or str(uuid4())
 39  
 40          output = {
 41              "question": chat_model.question,
 42              "type": "block",
 43              "sources": [],
 44              "guard": False,
 45              "tokens": {"input": 0, "output": 0},
 46              "project": project.props.name,
 47              "id": chat_id,
 48          }
 49  
 50          if self.check_input_guard(project, chat_model.question, user, db, output):
 51              yield output
 52              return
 53  
 54          # Load session and add user message
 55          session = await get_session(self.brain, chat_id)
 56          session.messages.append(user_text_message(chat_model.question))
 57  
 58          interpreter = BlockInterpreter(
 59              workspace_json=workspace,
 60              input_text=chat_model.question,
 61              brain=self.brain,
 62              user=user,
 63              db=db,
 64              image=chat_model.image,
 65              chat_id=chat_id,
 66              context=getattr(project, "context", None),
 67          )
 68          result = await interpreter.execute()
 69          output["answer"] = result or ""
 70  
 71          # Save assistant response to session
 72          session.messages.append(
 73              Message(role="assistant", content=[TextBlock(text=output["answer"])])
 74          )
 75          await save_session(self.brain, chat_id, session)
 76  
 77          yield output
 78  
 79      async def question(self, project: Project, question_model: QuestionModel, user: User, db: DBWrapper):
 80          workspace = self._get_workspace(project)
 81  
 82          interpreter = BlockInterpreter(
 83              workspace_json=workspace,
 84              input_text=question_model.question,
 85              brain=self.brain,
 86              user=user,
 87              db=db,
 88              image=question_model.image,
 89              context=getattr(project, "context", None),
 90          )
 91          result = await interpreter.execute()
 92  
 93          output = {
 94              "question": question_model.question,
 95              "type": "block",
 96              "sources": [],
 97              "guard": False,
 98              "tokens": {"input": 0, "output": 0},
 99              "project": project.props.name,
100              "answer": result,
101          }
102  
103          yield output