# OpenAI with Guardrail

#### **Use Case Overview**

We have created a chatbot, dividing it into two sides: the user side and the admin side. This is a simulation where the user interacts with the chatbot by asking questions, and the admin is the owner or the one who defines the criteria for the chatbot's questions.

Users can ask and answer questions just like the previous chatbot we created. However, there is a difference; they can only ask questions within the specified criteria. For example, this chatbot is set to only respond to questions related to animals. If the user asks a question unrelated to animals, the chatbot will reply that it cannot answer that question.

On the admin side, there is the privilege to set the criteria for the number of questions, and it is possible to add and clear criteria in the admin section of the chatbot, which displays all the criteria currently in use.

### User Flow

#### **Admin**

![Admin Interface for Rule Management](https://assets-global.website-files.com/63cb6b155c56b2dcd14e411d/656a2e7928457dee4ce1d53c_PO_3.png)

#### **User**

![User Interface for Question Submission](https://assets-global.website-files.com/63cb6b155c56b2dcd14e411d/656a2e80bc0ada798aebfc8a_PO_4.png)

### Frontend Development

#### **Admin Interface for Rule Management**

In designing the UI for the admin's rule management, the admin can create and delete rules and view the rules set for development. For the admin interface, React is used to manage various states. For form management, the library `react-hook-form`, `@hookform/resolvers/zod`, and `zod` are used for input validation. The connection to the API is handled using `@tanstack/react-query`, which provides various features for API interaction, such as using `useQueries` for fetching multiple values simultaneously.

```typescript
const [{ data: session }, { data: rule, refetch: refetchRule }] = useQueries({
    queries: [
      {
        queryKey: ['session', endPoint],
        queryFn: async () => {
          const { data } = await axios.get(`/session`)
          return data
        },
        enabled: typeof window !== 'undefined' && !localStorage?.session,
      },
      {
        queryKey: ['rule', endPoint],
        queryFn: async () => {
          const { data } = await axios.get(`/rule`)
          return data
        },
      },
    ],
  })
```

For adding and clearing rules, `useMutation` is used. The advantage of this library is that it handles various states automatically. For example, after successfully adding or clearing a rule, it triggers the `refetchRule` function to update the table.

```typescript
const { mutateAsync: createRule } = useMutation({
    mutationFn: (rule: string) => {
      return axios.post('/create_rule', { rule })
    },
    onSuccess: () => {
      refetchRule()
    },
  })

  const { mutateAsync: clearRule } = useMutation({
    mutationFn: () => {
      return axios.get('/clear_collection')
    },
    onSuccess: () => {
      refetchRule()
    },
  })
```

#### **User Interface for Question Submission**

In the implementation of the chatbot interface, there is an input section for typing messages to the AI. When a user types something unrelated to the specified rules, the API responds that it is a "bad prompt" and prompts the user to ask another question.

Now, let's see what happens when we ask a question related to cats. The AI responds by asking about the cat we inquired about. The interaction with the chatbot is done in a streaming format.

```typescript
useEffect(() => {
    const message = { query: watch('query') }
    const getData = async () => {
      try {
        setValue('query', '')
        const response = await fetch(
          `${endPoint ?? API_URL}/query?uuid=${localStorage?.session}&message=${
            message.query
          }`,
          {
            method: 'GET',
            headers: {
              Accept: 'text/event-stream',
              'x-api-key': localStorage?.apiKey,
            },
          }
        )

        if (response.status === 200) {
          const reader = response.body!.getReader()
          let result = ''
          while (true) {
            const { done, value } = await reader?.read()
            if (done) {
              setStreamText('')
              setAnswer((prevState) => [
                ...prevState,
                {
                  id: (prevState.length + 1).toString(),
                  role: 'ai',
                  message: result,
                },
              ])
              break
            }
            result += new TextDecoder().decode(value)
            setStreamText(result)
          }
        } else {
          if (response.status === 404) {
            localStorage.removeItem('session')
            refreshSession()
            setError('bot', {
              message: 'Something went wrong, please try again',
            })
            return
          }

          setError('bot', {
            message: 'Something went wrong',
          })
        }
      } catch (error: any) {
        console.error(error)
        setError('bot', {
          message: error?.response?.data?.message ?? 'Something went wrong',
        })
      }
    }
    if (isSubmitSuccessful && localStorage) {
      getData()
    }
  }, [
    submitCount,
    isSubmitSuccessful,
    setValue,
    watch,
    setError,
    endPoint,
    refreshSession,
  ])
```

#### **Guardrail Interaction**

The Guardrail feature limits the chatbot's ability to respond to prompts based on the rules set by the admin. The Relevant Answer Generation (RAG) checks how closely the prompt matches the rules and sends it to OpenAI. If the similarity is less than 99%, the API responds with a "bad prompt" message, indicating that the question is outside the defined rules.

#### **Error Handling and User Feedback**

Error handling is implemented for server connection issues or user input errors. For example, if there's a 404 status code (not found), the session is cleared and a new session is fetched. Validation errors for user input are handled using `react-hook-form`, `@hookform/resolvers/zod`, and `zod`.

```typescript
export const askScheme = z.object({
  apiKey: z.string().optional(),
  query: z.string().trim().min(1, { message: 'Please enter your message' }),
  bot: z.string({}).optional(),
})

export const ruleScheme = z.object({
  rule: z.string().trim().min(1, { message: 'Please enter your rule' }),
})

export const endpointScheme = z.object({
  endpoint: z.string().url().optional(),
})

const methodRule = useForm<IRuleForm>({
    resolver: zodResolver(ruleScheme),
    mode: 'onChange',
    shouldFocusError: true,
  })

  const methodsEndpoint = useForm<IEndPointForm>({
    resolver: zodResolver(endpointScheme),
    mode: 'onChange',
    shouldFocusError: true,
  })

  const methods = useForm<IOpenAIForm>({
    resolver: zodResolver(askScheme),
    mode: 'onChange',
    shouldFocusError: true,
    defaultValues: {
      apiKey: '',
      query: '',
    },
  })
```

[Link to Frontend Code](https://github.com/vultureprime/ai-web-interface/tree/main/next-open-ai-guardrail)

### Backend Development

#### **Setup FastAPI**

Install the necessary Python 3 libraries for FastAPI:

```python
pip install fastapi
pip install "uvicorn[standard]"
```

Create a file named `guardrail.py` and initialize FastAPI:

```python
from fastapi import FastAPI
from fastapi.encoders import jsonable_encoder
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/helloworld")
async def helloworld():
    return {"message": "Hello World"}
```

Run the FastAPI server using the command:

```python
uvicorn guardrail:app
```

This command instructs FastAPI, declared in the `app.py` file, to run, with the server defaulting to port 8000.

The API is divided into two parts: one for users sending prompts to the backend for interacting with AI, and another for administrators controlling whether prompts can be sent to AI.

Install other necessary dependencies:

```python
pip install openai
pip install uuid
pip install pydantic
pip install langchain
```

The admin API includes functions for creating, viewing, and deleting rules for filtering prompts, essentially adding data to VectorDB. The VectorDB initialization will be discussed further.

```python
def add_rule(text):
    Chroma.from_texts(collection_name=collection_name,texts=[text], embedding=embedding ,persist_directory=persist_dir)

@app.post('/create_rule')
def create_rule(rule:Rule):
    add_rule(rule.rule)
    return JSONResponse(content={
        "message":"success",
        "rule":rule.rule
        })
```

View all created rules:

```python
def get_collection():
    vectorstore = init_db()
    result = vectorstore.get()
    return result['documents']

@app.get('/rule')
def get_rule():
    res = get_collection()
    return JSONResponse(content={"result": res})
```

Clear all rules:

```python
def clear_db_collection():
    vectorstore = init_db()
    res = vectorstore.delete_collection()

@app.get('/clear_collection')
def clear_collection():
    clear_db_collection()
    return JSONResponse(content={"message": "remove complete"})
```

For user interaction, an API is provided to send prompts and check whether they match any rules:

```python
@app.get("/session")
def session():
    client_uuid = uuid.uuid4()
    create_session(str(client_uuid))
    result = {
        "uuid" : str(client_uuid)
    }
    return JSONResponse(content=result)

@app.get("/query")
async def main(uuid:str,message:str):
    get_session(uuid)
    collection = get_collection()
    if  len(collection) == 0:
        return StreamingResponse(
                    generate_response('No rules configure please ask admin'), 
                    media_type="text/event-stream"
                )
    else:
        score = compare_similarity(message)
        print(score)
        if score > 0.7:
            return StreamingResponse(
                        stream_chat(
                            uuid = uuid,
                            prompt= message
                        ), 
                        media_type="text/event-stream"
                    )
        else:
            return StreamingResponse(
                        generate_response('bad prompt please ask another one'), 
                        media_type="text/event-stream"
                    )
```

#### **Integrating OpenAI API**

Connect to the OpenAI API using the API key generated from the OpenAI console. Two methods are provided: one using an Embedding model and the other using the Chatbot API.

1. Embedding model:

```python
from langchain.embeddings import OpenAIEmbeddings
api_key = 'sk-XcTnjgYVsJQMNxxxxxxxxxxxxxx'
embedding = OpenAIEmbeddings(openai_api_key=api_key)
```

2. Chatbot API:

```python
import openai
openai.api_key = api_key

def stream_chat(uuid: str, prompt: str):
    result = ""
    messages = add_message(uuid, 'user', prompt)
    for chunk in openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=messages,
        stream=True,
    ):
        content = chunk["choices"][0].get("delta", {}).get("content")
        if content is not None:
            result = result + content
            yield content
    add_message(uuid, 'assistant', result)
```

#### **Implementing VectorDB with Chroma**

Chroma is used for VectorDB, which stores vectorized data to aid in finding similar data. Initialization is done using parameters for collection name, data location, and the embedding model.

```python
from langchain.vectorstores import Chroma

def init_db():
    vectorstore = Chroma(collection_name=collection_name, persist_directory=persist_dir, embedding_function=embedding)
    return vectorstore
```

#### **Building the Guardrail Feature**

The Guardrail feature queries rule data from VectorDB based on received prompts, and if the similarity score is above a threshold, the prompt is forwarded to the chatbot.

```python
def compare_similarity(query):
    vectorstore = init_db()
    result = vectorstore.similarity_search_with_relevance_scores(query, k=5)
    score_list = []
    for i in result:
        score_list.append(i[-1])
    try:
        average = sum(score_list)/len(score_list)
        return average
    except:
        return 0
```

#### **Deploying on AWS EC2**

Deployment is done using the `screen` utility:

1. Create a session with a specific name:

```bash
screen -S name
```

2. Navigate to the API folder:

```bash
cd path/to/api
```

3. Start FastAPI on port 8000:

```bash
uvicorn guardrail:app --host 0.0.0.0 --port 8000
```

4. Detach from the current screen session:

```bash
Ctrl+a d
```

#### **API Gateway and CORS Configuration**

Add CORS configuration for FastAPI:

```python
from fastapi.middleware.cors import CORSMiddleware

app.app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
```

Connect this server to AWS API Gateway for authentication and usage management in each request.

[Link to Backend Code](https://github.com/vultureprime/ai-web-backend/tree/main/Guardrail-openai)